Added support for async-std test attribute using the inventory crate

Author: Andrew J Westlake

Cargo.toml

@@ -42,6 +42,13 @@ name = "tokio_multi_thread"
 path = "examples/tokio_multi_thread.rs"
 required-features = ["attributes", "tokio-runtime"]
 
+
+[[test]]
+name = "test_async_std_test_macro"
+path = "pytests/test_async_std_test_macro.rs"
+harness = false
+required-features = ["attributes", "async-std-runtime", "testing"]
+
 [[test]]
 name = "test_async_std_asyncio"
 path = "pytests/test_async_std_asyncio.rs"
@@ -94,4 +101,7 @@ optional = true
 [dependencies.tokio]
 version = "1.0" 
 features = ["full"]
-optional = true
+optional = true
+
+[dev-dependencies]
+inventory = "0.1"

pyo3-asyncio-macros/src/lib.rs

@@ -102,3 +102,93 @@ pub fn async_std_main(_attr: TokenStream, item: TokenStream) -> TokenStream {
 pub fn tokio_main(args: TokenStream, item: TokenStream) -> TokenStream {
     tokio::main(args, item, true)
 }
+
+/// Enables an async main function that uses the async-std runtime.
+///
+/// # Examples
+///
+/// ```ignore
+/// #[pyo3_asyncio::async_std::main]
+/// async fn main() -> PyResult<()> {
+///     Ok(())
+/// }
+/// ```
+#[cfg(not(test))] // NOTE: exporting main breaks tests, we should file an issue.
+#[proc_macro_attribute]
+pub fn async_std_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = syn::parse_macro_input!(item as syn::ItemFn);
+
+    let sig = &input.sig;
+    // let inputs = &input.sig.inputs;
+    let name = &input.sig.ident;
+    let body = &input.block;
+    // let attrs = &input.attrs;
+    let vis = &input.vis;
+
+    let fn_impl = if input.sig.asyncness.is_none() {
+        quote! {
+            #vis fn #name() -> std::pin::Pin<Box<dyn std::future::Future<Output = pyo3::PyResult<()>> + Send>> {
+                #sig {
+                    #body
+                }
+
+                Box::pin(async_std::task::spawn_blocking(move || {
+                    #name();
+                    Ok(())
+                }))
+            }
+        }
+    } else {
+        quote! {
+            #vis fn #name() -> std::pin::Pin<Box<dyn std::future::Future<Output = pyo3::PyResult<()>> + Send>> {
+                #sig {
+                    #body
+                }
+
+                Box::pin(#name())
+            }
+        }
+    };
+
+    let result = quote! {
+        #fn_impl
+
+        inventory::submit!(crate::Test { name: stringify!(#name), test_fn: &#name });
+    };
+
+    result.into()
+}
+
+#[cfg(not(test))]
+#[proc_macro]
+pub fn async_std_test_main(args: TokenStream) -> TokenStream {
+    let suite_name = syn::parse_macro_input!(args as syn::LitStr);
+
+    let result = quote! {
+        #[derive(Clone)]
+        pub(crate) struct Test {
+            pub name: &'static str,
+            pub test_fn: &'static (dyn Fn() -> std::pin::Pin<Box<dyn std::future::Future<Output = pyo3::PyResult<()>> + Send>> + Send + Sync),
+        }
+
+        impl pyo3_asyncio::testing::TestTrait for Test {
+            fn name(&self) -> &str {
+                self.name
+            }
+
+            fn task(self) -> std::pin::Pin<Box<dyn std::future::Future<Output = pyo3::PyResult<()>> + Send>> {
+                (self.test_fn)()
+            }
+        }
+
+        inventory::collect!(Test);
+
+        fn main() {
+            pyo3_asyncio::async_std::testing::test_main(
+                #suite_name,
+                inventory::iter::<Test>().map(|test| test.clone()).collect()
+            );
+        }
+    };
+    result.into()
+}

pytests/test_async_std_test_macro.rs

@@ -0,0 +1,17 @@
+use std::{time::Duration, thread};
+
+use pyo3::prelude::*;
+
+#[pyo3_asyncio::async_std::test]
+fn test() {
+    thread::sleep(Duration::from_secs(1));
+}
+
+#[pyo3_asyncio::async_std::test]
+async fn test_sleep() -> PyResult<()> {
+    async_std::task::sleep(Duration::from_secs(1)).await;
+
+    Ok(())
+}
+
+pyo3_asyncio::async_std::test_main!("PyO3 test async-std test macros");

src/async_std.rs

@@ -8,6 +8,11 @@ use crate::generic::{self, JoinError, Runtime};
 /// <span class="module-item stab portability" style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"><code>attributes</code></span> Sets up the async-std runtime and runs an async fn as main
 #[cfg(feature = "attributes")]
 pub use pyo3_asyncio_macros::async_std_main as main;
+#[cfg(feature = "attributes")]
+pub use pyo3_asyncio_macros::async_std_test as test;
+
+#[cfg(feature = "attributes")]
+pub use pyo3_asyncio_macros::async_std_test_main as test_main;
 
 struct AsyncStdJoinError;
 
@@ -157,7 +162,10 @@ pub mod testing {
     //!
     //! ```no_run
     //! fn main() {
-    //!     pyo3_asyncio::async_std::testing::test_main("Example Test Suite", vec![]);
+    //!     pyo3_asyncio::async_std::testing::test_main(
+    //!         "Example Test Suite",
+    //!         Vec::<pyo3_asyncio::testing::Test>::new()
+    //!     );
     //! }
     //! ```
     //!
@@ -194,7 +202,11 @@ pub mod testing {
     use async_std::task;
     use pyo3::prelude::*;
 
-    use crate::{async_std::AsyncStdRuntime, generic, testing::Test};
+    use crate::{
+        async_std::AsyncStdRuntime,
+        generic,
+        testing::{Test, TestTrait},
+    };
 
     /// Construct a test from a blocking function (like the traditional `#[test]` attribute)
     pub fn new_sync_test<F>(name: String, func: F) -> Test
@@ -208,7 +220,7 @@ pub mod testing {
     ///
     /// This is meant to perform the necessary initialization for most test cases. If you want
     /// additional control over the initialization, you can use this function as a template.
-    pub fn test_main(suite_name: &str, tests: Vec<Test>) {
-        generic::testing::test_main::<AsyncStdRuntime>(suite_name, tests)
+    pub fn test_main(suite_name: &str, tests: Vec<impl TestTrait + 'static>) {
+        generic::testing::test_main::<AsyncStdRuntime, _>(suite_name, tests)
     }
 }

src/generic.rs

@@ -244,7 +244,7 @@ pub mod testing {
     use crate::{
         dump_err,
         generic::{run_until_complete, Runtime},
-        testing::{parse_args, test_harness, Test},
+        testing::{parse_args, test_harness, TestTrait},
         with_runtime,
     };
 
@@ -253,9 +253,10 @@ pub mod testing {
     /// This is meant to perform the necessary initialization for most test cases. If you want
     /// additional control over the initialization, you can use this
     /// function as a template.
-    pub fn test_main<R>(suite_name: &str, tests: Vec<Test>)
+    pub fn test_main<R, T>(suite_name: &str, tests: Vec<T>)
     where
         R: Runtime,
+        T: TestTrait + 'static,
     {
         Python::with_gil(|py| {
             with_runtime(py, || {

src/testing.rs

@@ -82,6 +82,20 @@ pub struct Test {
     task: Pin<Box<dyn Future<Output = PyResult<()>> + Send>>,
 }
 
+pub trait TestTrait: Send {
+    fn name(&self) -> &str;
+    fn task(self) -> Pin<Box<dyn Future<Output = PyResult<()>> + Send>>;
+}
+
+impl TestTrait for Test {
+    fn name(&self) -> &str {
+        self.name.as_str()
+    }
+    fn task(self) -> Pin<Box<dyn Future<Output = PyResult<()>> + Send>> {
+        self.task
+    }
+}
+
 impl Test {
     /// Construct a test from a future
     pub fn new_async(
@@ -96,21 +110,23 @@ impl Test {
 }
 
 /// Run a sequence of tests while applying any necessary filtering from the `Args`
-pub async fn test_harness(tests: Vec<Test>, args: Args) -> PyResult<()> {
+pub async fn test_harness(tests: Vec<impl TestTrait + 'static>, args: Args) -> PyResult<()> {
     stream::iter(tests)
         .for_each_concurrent(Some(4), |test| {
             let mut ignore = false;
 
             if let Some(filter) = args.filter.as_ref() {
-                if !test.name.contains(filter) {
+                if !test.name().contains(filter) {
                     ignore = true;
                 }
             }
 
             async move {
                 if !ignore {
-                    test.task.await.unwrap();
-                    println!("test {} ... ok", test.name);
+                    let name = test.name().to_string();
+                    test.task().await.unwrap();
+
+                    println!("test {} ... ok", name);
                 }
             }
         })