Repository: tokio-rs/tokio (language: rust)

Issue:
Allow setting `unhandled_panic` behavior as option on `tokio::test`
**Is your feature request related to a problem? Please describe.**
I have several unit tests that run some handler code that is under test in a `tokio::spawn`ed task, and sends/receives bytes to/from that handler code from the main task. My AsyncRead + AsyncWrite mock will panic if it sees unexpected bytes, and if this happens in the background task the test will hang. I'd prefer the test to shut down in this scenario, and so I'm using the `unhandled_panic` option introduced by #4516.

**Describe the solution you'd like**
`#[tokio::test(unhandled_panic = ShutdownRuntime)`

**Describe alternatives you've considered**
Currently I manually set up a tokio runtime for my tests that require this behavior.



Discussion / hints:
I’m fine adding this feature. 

Fix patch:
diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs
index 8858c8a1674..acdc2610f44 100644
--- a/tokio-macros/src/entry.rs
+++ b/tokio-macros/src/entry.rs
@@ -25,11 +25,37 @@ impl RuntimeFlavor {
     }
 }
 
+#[derive(Clone, Copy, PartialEq)]
+enum UnhandledPanic {
+    Ignore,
+    ShutdownRuntime,
+}
+
+impl UnhandledPanic {
+    fn from_str(s: &str) -> Result<UnhandledPanic, String> {
+        match s {
+            "ignore" => Ok(UnhandledPanic::Ignore),
+            "shutdown_runtime" => Ok(UnhandledPanic::ShutdownRuntime),
+            _ => Err(format!("No such unhandled panic behavior `{}`. The unhandled panic behaviors are `ignore` and `shutdown_runtime`.", s)),
+        }
+    }
+
+    fn into_tokens(self, crate_path: &TokenStream) -> TokenStream {
+        match self {
+            UnhandledPanic::Ignore => quote! { #crate_path::runtime::UnhandledPanic::Ignore },
+            UnhandledPanic::ShutdownRuntime => {
+                quote! { #crate_path::runtime::UnhandledPanic::ShutdownRuntime }
+            }
+        }
+    }
+}
+
 struct FinalConfig {
     flavor: RuntimeFlavor,
     worker_threads: Option<usize>,
     start_paused: Option<bool>,
     crate_name: Option<Path>,
+    unhandled_panic: Option<UnhandledPanic>,
 }
 
 /// Config used in case of the attribute not being able to build a valid config
@@ -38,6 +64,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
     worker_threads: None,
     start_paused: None,
     crate_name: None,
+    unhandled_panic: None,
 };
 
 struct Configuration {
@@ -48,6 +75,7 @@ struct Configuration {
     start_paused: Option<(bool, Span)>,
     is_test: bool,
     crate_name: Option<Path>,
+    unhandled_panic: Option<(UnhandledPanic, Span)>,
 }
 
 impl Configuration {
@@ -63,6 +91,7 @@ impl Configuration {
             start_paused: None,
             is_test,
             crate_name: None,
+            unhandled_panic: None,
         }
     }
 
@@ -117,6 +146,25 @@ impl Configuration {
         Ok(())
     }
 
+    fn set_unhandled_panic(
+        &mut self,
+        unhandled_panic: syn::Lit,
+        span: Span,
+    ) -> Result<(), syn::Error> {
+        if self.unhandled_panic.is_some() {
+            return Err(syn::Error::new(
+                span,
+                "`unhandled_panic` set multiple times.",
+            ));
+        }
+
+        let unhandled_panic = parse_string(unhandled_panic, span, "unhandled_panic")?;
+        let unhandled_panic =
+            UnhandledPanic::from_str(&unhandled_panic).map_err(|err| syn::Error::new(span, err))?;
+        self.unhandled_panic = Some((unhandled_panic, span));
+        Ok(())
+    }
+
     fn macro_name(&self) -> &'static str {
         if self.is_test {
             "tokio::test"
@@ -163,11 +211,24 @@ impl Configuration {
             (_, None) => None,
         };
 
+        let unhandled_panic = match (flavor, self.unhandled_panic) {
+            (F::Threaded, Some((_, unhandled_panic_span))) => {
+                let msg = format!(
+                    "The `unhandled_panic` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
+                    self.macro_name(),
+                );
+                return Err(syn::Error::new(unhandled_panic_span, msg));
+            }
+            (F::CurrentThread, Some((unhandled_panic, _))) => Some(unhandled_panic),
+            (_, None) => None,
+        };
+
         Ok(FinalConfig {
             crate_name: self.crate_name.clone(),
             flavor,
             worker_threads,
             start_paused,
+            unhandled_panic,
         })
     }
 }
@@ -275,9 +336,13 @@ fn build_config(
                     "crate" => {
                         config.set_crate_name(lit.clone(), syn::spanned::Spanned::span(lit))?;
                     }
+                    "unhandled_panic" => {
+                        config
+                            .set_unhandled_panic(lit.clone(), syn::spanned::Spanned::span(lit))?;
+                    }
                     name => {
                         let msg = format!(
-                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`",
+                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`",
                             name,
                         );
                         return Err(syn::Error::new_spanned(namevalue, msg));
@@ -303,11 +368,11 @@ fn build_config(
                             macro_name
                         )
                     }
-                    "flavor" | "worker_threads" | "start_paused" => {
+                    "flavor" | "worker_threads" | "start_paused" | "crate" | "unhandled_panic" => {
                         format!("The `{}` attribute requires an argument.", name)
                     }
                     name => {
-                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`", name)
+                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.", name)
                     }
                 };
                 return Err(syn::Error::new_spanned(path, msg));
@@ -359,6 +424,10 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
     if let Some(v) = config.start_paused {
         rt = quote_spanned! {last_stmt_start_span=> #rt.start_paused(#v) };
     }
+    if let Some(v) = config.unhandled_panic {
+        let unhandled_panic = v.into_tokens(&crate_path);
+        rt = quote_spanned! {last_stmt_start_span=> #rt.unhandled_panic(#unhandled_panic) };
+    }
 
     let generated_attrs = if is_test {
         quote! {
diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs
index c108d8c46a2..860a0929952 100644
--- a/tokio-macros/src/lib.rs
+++ b/tokio-macros/src/lib.rs
@@ -202,6 +202,52 @@ use proc_macro::TokenStream;
 ///         })
 /// }
 /// ```
+///
+/// ### Configure unhandled panic behavior
+///
+/// Available options are `shutdown_runtime` and `ignore`. For more details, see
+/// [`Builder::unhandled_panic`].
+///
+/// This option is only compatible with the `current_thread` runtime.
+///
+/// ```no_run
+/// #[cfg(tokio_unstable)]
+/// #[tokio::main(flavor = "current_thread", unhandled_panic = "shutdown_runtime")]
+/// async fn main() {
+///     let _ = tokio::spawn(async {
+///         panic!("This panic will shutdown the runtime.");
+///     }).await;
+/// }
+/// # #[cfg(not(tokio_unstable))]
+/// # fn main() { }
+/// ```
+///
+/// Equivalent code not using `#[tokio::main]`
+///
+/// ```no_run
+/// #[cfg(tokio_unstable)]
+/// fn main() {
+///     tokio::runtime::Builder::new_current_thread()
+///         .enable_all()
+///         .unhandled_panic(UnhandledPanic::ShutdownRuntime)
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             let _ = tokio::spawn(async {
+///                 panic!("This panic will shutdown the runtime.");
+///             }).await;
+///         })
+/// }
+/// # #[cfg(not(tokio_unstable))]
+/// # fn main() { }
+/// ```
+///
+/// **Note**: This option depends on Tokio's [unstable API][unstable]. See [the
+/// documentation on unstable features][unstable] for details on how to enable
+/// Tokio's unstable features.
+///
+/// [`Builder::unhandled_panic`]: ../tokio/runtime/struct.Builder.html#method.unhandled_panic
+/// [unstable]: ../tokio/index.html#unstable-features
 #[proc_macro_attribute]
 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::main(args.into(), item.into(), true).into()
@@ -423,6 +469,53 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
 ///     println!("Hello world");
 /// }
 /// ```
+///
+/// ### Configure unhandled panic behavior
+///
+/// Available options are `shutdown_runtime` and `ignore`. For more details, see
+/// [`Builder::unhandled_panic`].
+///
+/// This option is only compatible with the `current_thread` runtime.
+///
+/// ```no_run
+/// #[cfg(tokio_unstable)]
+/// #[tokio::test(flavor = "current_thread", unhandled_panic = "shutdown_runtime")]
+/// async fn my_test() {
+///     let _ = tokio::spawn(async {
+///         panic!("This panic will shutdown the runtime.");
+///     }).await;
+/// }
+/// # #[cfg(not(tokio_unstable))]
+/// # fn main() { }
+/// ```
+///
+/// Equivalent code not using `#[tokio::test]`
+///
+/// ```no_run
+/// #[cfg(tokio_unstable)]
+/// #[test]
+/// fn my_test() {
+///     tokio::runtime::Builder::new_current_thread()
+///         .enable_all()
+///         .unhandled_panic(UnhandledPanic::ShutdownRuntime)
+///         .build()
+///         .unwrap()
+///         .block_on(async {
+///             let _ = tokio::spawn(async {
+///                 panic!("This panic will shutdown the runtime.");
+///             }).await;
+///         })
+/// }
+/// # #[cfg(not(tokio_unstable))]
+/// # fn main() { }
+/// ```
+///
+/// **Note**: This option depends on Tokio's [unstable API][unstable]. See [the
+/// documentation on unstable features][unstable] for details on how to enable
+/// Tokio's unstable features.
+///
+/// [`Builder::unhandled_panic`]: ../tokio/runtime/struct.Builder.html#method.unhandled_panic
+/// [unstable]: ../tokio/index.html#unstable-features
 #[proc_macro_attribute]
 pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
     entry::test(args.into(), item.into(), true).into()


Test patch:
diff --git a/tests-build/tests/fail/macros_invalid_input.rs b/tests-build/tests/fail/macros_invalid_input.rs
index 85b4924f225..425b274ff7f 100644
--- a/tests-build/tests/fail/macros_invalid_input.rs
+++ b/tests-build/tests/fail/macros_invalid_input.rs
@@ -41,6 +41,9 @@ async fn test_crate_not_path_int() {}
 #[tokio::test(crate = "456")]
 async fn test_crate_not_path_invalid() {}
 
+#[tokio::test(flavor = "multi_thread", unhandled_panic = "shutdown_runtime")]
+async fn test_multi_thread_with_unhandled_panic() {}
+
 #[tokio::test]
 #[test]
 async fn test_has_second_test_attr() {}
diff --git a/tests-build/tests/fail/macros_invalid_input.stderr b/tests-build/tests/fail/macros_invalid_input.stderr
index 11f95315cdf..e6e80e8114a 100644
--- a/tests-build/tests/fail/macros_invalid_input.stderr
+++ b/tests-build/tests/fail/macros_invalid_input.stderr
@@ -1,115 +1,121 @@
 error: the `async` keyword is missing from the function declaration
- --> $DIR/macros_invalid_input.rs:6:1
+ --> tests/fail/macros_invalid_input.rs:6:1
   |
 6 | fn main_is_not_async() {}
   | ^^
 
-error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
- --> $DIR/macros_invalid_input.rs:8:15
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.
+ --> tests/fail/macros_invalid_input.rs:8:15
   |
 8 | #[tokio::main(foo)]
   |               ^^^
 
 error: Must have specified ident
-  --> $DIR/macros_invalid_input.rs:11:15
+  --> tests/fail/macros_invalid_input.rs:11:15
    |
 11 | #[tokio::main(threadpool::bar)]
    |               ^^^^^^^^^^^^^^^
 
 error: the `async` keyword is missing from the function declaration
-  --> $DIR/macros_invalid_input.rs:15:1
+  --> tests/fail/macros_invalid_input.rs:15:1
    |
 15 | fn test_is_not_async() {}
    | ^^
 
-error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
-  --> $DIR/macros_invalid_input.rs:17:15
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.
+  --> tests/fail/macros_invalid_input.rs:17:15
    |
 17 | #[tokio::test(foo)]
    |               ^^^
 
-error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
-  --> $DIR/macros_invalid_input.rs:20:15
+error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`
+  --> tests/fail/macros_invalid_input.rs:20:15
    |
 20 | #[tokio::test(foo = 123)]
    |               ^^^^^^^^^
 
 error: Failed to parse value of `flavor` as string.
-  --> $DIR/macros_invalid_input.rs:23:24
+  --> tests/fail/macros_invalid_input.rs:23:24
    |
 23 | #[tokio::test(flavor = 123)]
    |                        ^^^
 
 error: No such runtime flavor `foo`. The runtime flavors are `current_thread` and `multi_thread`.
-  --> $DIR/macros_invalid_input.rs:26:24
+  --> tests/fail/macros_invalid_input.rs:26:24
    |
 26 | #[tokio::test(flavor = "foo")]
    |                        ^^^^^
 
 error: The `start_paused` option requires the `current_thread` runtime flavor. Use `#[tokio::test(flavor = "current_thread")]`
-  --> $DIR/macros_invalid_input.rs:29:55
+  --> tests/fail/macros_invalid_input.rs:29:55
    |
 29 | #[tokio::test(flavor = "multi_thread", start_paused = false)]
    |                                                       ^^^^^
 
 error: Failed to parse value of `worker_threads` as integer.
-  --> $DIR/macros_invalid_input.rs:32:57
+  --> tests/fail/macros_invalid_input.rs:32:57
    |
 32 | #[tokio::test(flavor = "multi_thread", worker_threads = "foo")]
    |                                                         ^^^^^
 
 error: The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[tokio::test(flavor = "multi_thread")]`
-  --> $DIR/macros_invalid_input.rs:35:59
+  --> tests/fail/macros_invalid_input.rs:35:59
    |
 35 | #[tokio::test(flavor = "current_thread", worker_threads = 4)]
    |                                                           ^
 
 error: Failed to parse value of `crate` as path.
-  --> $DIR/macros_invalid_input.rs:38:23
+  --> tests/fail/macros_invalid_input.rs:38:23
    |
 38 | #[tokio::test(crate = 456)]
    |                       ^^^
 
 error: Failed to parse value of `crate` as path: "456"
-  --> $DIR/macros_invalid_input.rs:41:23
+  --> tests/fail/macros_invalid_input.rs:41:23
    |
 41 | #[tokio::test(crate = "456")]
    |                       ^^^^^
 
+error: The `unhandled_panic` option requires the `current_thread` runtime flavor. Use `#[tokio::test(flavor = "current_thread")]`
+  --> tests/fail/macros_invalid_input.rs:44:58
+   |
+44 | #[tokio::test(flavor = "multi_thread", unhandled_panic = "shutdown_runtime")]
+   |                                                          ^^^^^^^^^^^^^^^^^^
+
 error: second test attribute is supplied, consider removing or changing the order of your test attributes
-  --> $DIR/macros_invalid_input.rs:45:1
+  --> tests/fail/macros_invalid_input.rs:48:1
    |
-45 | #[test]
+48 | #[test]
    | ^^^^^^^
 
 error: second test attribute is supplied, consider removing or changing the order of your test attributes
-  --> $DIR/macros_invalid_input.rs:49:1
+  --> tests/fail/macros_invalid_input.rs:52:1
    |
-49 | #[::core::prelude::v1::test]
+52 | #[::core::prelude::v1::test]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 error: second test attribute is supplied, consider removing or changing the order of your test attributes
-  --> $DIR/macros_invalid_input.rs:53:1
+  --> tests/fail/macros_invalid_input.rs:56:1
    |
-53 | #[core::prelude::rust_2015::test]
+56 | #[core::prelude::rust_2015::test]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 error: second test attribute is supplied, consider removing or changing the order of your test attributes
-  --> $DIR/macros_invalid_input.rs:57:1
+  --> tests/fail/macros_invalid_input.rs:60:1
    |
-57 | #[::std::prelude::rust_2018::test]
+60 | #[::std::prelude::rust_2018::test]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 error: second test attribute is supplied, consider removing or changing the order of your test attributes
-  --> $DIR/macros_invalid_input.rs:61:1
+  --> tests/fail/macros_invalid_input.rs:64:1
    |
-61 | #[std::prelude::rust_2021::test]
+64 | #[std::prelude::rust_2021::test]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 error: second test attribute is supplied, consider removing or changing the order of your test attributes
-  --> $DIR/macros_invalid_input.rs:64:1
+  --> tests/fail/macros_invalid_input.rs:67:1
    |
-64 | #[tokio::test]
+67 | #[tokio::test]
    | ^^^^^^^^^^^^^^
    |
    = note: this error originates in the attribute macro `tokio::test` (in Nightly builds, run with -Z macro-backtrace for more info)
diff --git a/tokio/tests/macros_test.rs b/tokio/tests/macros_test.rs
index 69ee30b36ce..bed443cf293 100644
--- a/tokio/tests/macros_test.rs
+++ b/tokio/tests/macros_test.rs
@@ -1,3 +1,4 @@
+#![allow(unknown_lints, unexpected_cfgs)]
 #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threading
 
 use tokio::test;
@@ -92,3 +93,26 @@ pub mod issue_5243 {
         async fn foo() {}
     );
 }
+
+#[cfg(tokio_unstable)]
+pub mod macro_rt_arg_unhandled_panic {
+    use tokio_test::assert_err;
+
+    #[tokio::test(flavor = "current_thread", unhandled_panic = "shutdown_runtime")]
+    #[should_panic]
+    async fn unhandled_panic_shutdown_runtime() {
+        let _ = tokio::spawn(async {
+            panic!("This panic should shutdown the runtime.");
+        })
+        .await;
+    }
+
+    #[tokio::test(flavor = "current_thread", unhandled_panic = "ignore")]
+    async fn unhandled_panic_ignore() {
+        let rt = tokio::spawn(async {
+            panic!("This panic should be forwarded to rt as an error.");
+        })
+        .await;
+        assert_err!(rt);
+    }
+}