From c78c9168d7d9e7dce4b3ab5426a68a9787bae147 Mon Sep 17 00:00:00 2001
From: Jon Gjengset <jon@thesquareplanet.com>
Date: Wed, 2 Oct 2019 13:58:34 -0400
Subject: [PATCH] macros: allow selecting runtime in tokio::test attr (#1620)

In the past, it was not possible to choose to use the multi-threaded
tokio `Runtime` in tests, which meant that any test that transitively
used `executor::threadpool::blocking` would fail with

```
'blocking' annotation used from outside the context of a thread pool
```

This patch adds a runtime annotation attribute to `#[tokio::test]` just
like `#[tokio::main]` has, which lets users opt in to the threadpool
runtime over `current_thread` (the default).
---
 .../tests/fail/macros_invalid_input.stderr    |  4 +-
 tokio-macros/src/lib.rs                       | 74 +++++++++++++++----
 2 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/build-tests/tests/fail/macros_invalid_input.stderr b/build-tests/tests/fail/macros_invalid_input.stderr
index 7aa9049905c..14061e3520a 100644
--- a/build-tests/tests/fail/macros_invalid_input.stderr
+++ b/build-tests/tests/fail/macros_invalid_input.stderr
@@ -4,7 +4,7 @@ error: the async keyword is missing from the function declaration
 4 | fn main_is_not_async() {}
   | ^^
 
-error: Unknown attribute foo is specified
+error: Unknown attribute foo is specified; expected `single_thread` or `multi_thread`
  --> $DIR/macros_invalid_input.rs:6:15
   |
 6 | #[tokio::main(foo)]
@@ -28,7 +28,7 @@ error: the test function cannot accept arguments
 16 | async fn test_fn_has_args(_x: u8) {}
    |                           ^^^^^^
 
-error: unexpected token
+error: Unknown attribute foo is specified; expected `single_thread` or `multi_thread`
   --> $DIR/macros_invalid_input.rs:18:15
    |
 18 | #[tokio::test(foo)]
diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs
index 07adfa0da4a..24a520213cb 100644
--- a/tokio-macros/src/lib.rs
+++ b/tokio-macros/src/lib.rs
@@ -18,6 +18,12 @@ extern crate proc_macro;
 use proc_macro::TokenStream;
 use quote::quote;
 
+enum RuntimeType {
+    Single,
+    Multi,
+    Auto,
+}
+
 /// Marks async function to be executed by selected runtime.
 ///
 /// ## Options:
@@ -50,12 +56,6 @@ use quote::quote;
 #[proc_macro_attribute]
 #[cfg(not(test))] // Work around for rust-lang/rust#62127
 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
-    enum RuntimeType {
-        Single,
-        Multi,
-        Auto,
-    }
-
     let input = syn::parse_macro_input!(item as syn::ItemFn);
     let args = syn::parse_macro_input!(args as syn::AttributeArgs);
 
@@ -90,7 +90,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
                 "multi_thread" => runtime = RuntimeType::Multi,
                 "single_thread" => runtime = RuntimeType::Single,
                 name => {
-                    let msg = format!("Unknown attribute {} is specified", name);
+                    let msg = format!("Unknown attribute {} is specified; expected `single_thread` or `multi_thread`", name);
                     return syn::Error::new_spanned(path, msg).to_compile_error().into();
                 }
             }
@@ -124,9 +124,23 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 
 /// Marks async function to be executed by runtime, suitable to test enviornment
 ///
-/// Uses `current_thread` runtime.
+/// ## Options:
+///
+/// - `single_thread` - Uses `current_thread`. Used by default.
+/// - `multi_thread` - Uses multi-threaded runtime.
+///
+/// ## Usage
+///
+/// ### Select runtime
+///
+/// ```no_run
+/// #[tokio::test(multi_thread)]
+/// async fn my_test() {
+///     assert!(true);
+/// }
+/// ```
 ///
-/// # Examples
+/// ### Using default
 ///
 /// ```no_run
 /// #[tokio::test]
@@ -137,7 +151,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 #[proc_macro_attribute]
 pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
     let input = syn::parse_macro_input!(item as syn::ItemFn);
-    let _ = syn::parse_macro_input!(args as syn::parse::Nothing);
+    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
 
     let ret = &input.sig.output;
     let name = &input.sig.ident;
@@ -165,13 +179,41 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
             .into();
     }
 
-    let result = quote! {
-        #[test]
-        #(#attrs)*
-        fn #name() #ret {
-            let mut rt = tokio::runtime::current_thread::Runtime::new().unwrap();
-            rt.block_on(async { #body })
+    let mut runtime = RuntimeType::Auto;
+
+    for arg in args {
+        if let syn::NestedMeta::Meta(syn::Meta::Path(path)) = arg {
+            let ident = path.get_ident();
+            if ident.is_none() {
+                let msg = "Must have specified ident";
+                return syn::Error::new_spanned(path, msg).to_compile_error().into();
+            }
+            match ident.unwrap().to_string().to_lowercase().as_str() {
+                "multi_thread" => runtime = RuntimeType::Multi,
+                "single_thread" => runtime = RuntimeType::Single,
+                name => {
+                    let msg = format!("Unknown attribute {} is specified; expected `single_thread` or `multi_thread`", name);
+                    return syn::Error::new_spanned(path, msg).to_compile_error().into();
+                }
+            }
         }
+    }
+
+    let result = match runtime {
+        RuntimeType::Multi => quote! {
+            #[test]
+            #(#attrs)*
+            fn #name() #ret {
+                tokio::runtime::Runtime::new().unwrap().block_on(async { #body })
+            }
+        },
+        RuntimeType::Single | RuntimeType::Auto => quote! {
+            #[test]
+            #(#attrs)*
+            fn #name() #ret {
+                tokio::runtime::current_thread::Runtime::new().unwrap().block_on(async { #body })
+            }
+        },
     };
 
     result.into()