Added tokio test attributes

Author: Andrew J Westlake

Cargo.toml

@@ -43,12 +43,6 @@ 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"

pyo3-asyncio-macros/src/lib.rs

@@ -103,26 +103,14 @@ 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() {
@@ -133,8 +121,7 @@ pub fn async_std_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
                 }
 
                 Box::pin(async_std::task::spawn_blocking(move || {
-                    #name();
-                    Ok(())
+                    #name()
                 }))
             }
         }
@@ -153,7 +140,10 @@ pub fn async_std_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
     let result = quote! {
         #fn_impl
 
-        inventory::submit!(crate::Test { name: stringify!(#name), test_fn: &#name });
+        inventory::submit!(crate::Test {
+            name: format!("{}::{}", std::module_path!(), stringify!(#name)),
+            test_fn: &#name
+        });
     };
 
     result.into()
@@ -167,13 +157,13 @@ pub fn async_std_test_main(args: TokenStream) -> TokenStream {
     let result = quote! {
         #[derive(Clone)]
         pub(crate) struct Test {
-            pub name: &'static str,
+            pub name: String,
             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
+                self.name.as_str()
             }
 
             fn task(self) -> std::pin::Pin<Box<dyn std::future::Future<Output = pyo3::PyResult<()>> + Send>> {
@@ -192,3 +182,90 @@ pub fn async_std_test_main(args: TokenStream) -> TokenStream {
     };
     result.into()
 }
+
+#[cfg(not(test))] // NOTE: exporting main breaks tests, we should file an issue.
+#[proc_macro_attribute]
+pub fn tokio_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = syn::parse_macro_input!(item as syn::ItemFn);
+
+    let sig = &input.sig;
+    let name = &input.sig.ident;
+    let body = &input.block;
+    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 {
+                    match pyo3_asyncio::tokio::get_runtime().spawn_blocking(&#name).await {
+                        Ok(result) => result,
+                        Err(e) => {
+                            assert!(e.is_panic());
+                            Err(pyo3::exceptions::PyException::new_err("rust future panicked"))
+                        }
+                    }
+                })
+            }
+        }
+    } 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: format!("{}::{}", std::module_path!(), stringify!(#name)),
+            test_fn: &#name
+        });
+    };
+
+    result.into()
+}
+
+#[cfg(not(test))]
+#[proc_macro]
+pub fn tokio_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: String,
+            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.as_str()
+            }
+
+            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::tokio::init_multi_thread();
+            pyo3_asyncio::tokio::testing::test_main(
+                #suite_name,
+                inventory::iter::<Test>().map(|test| test.clone()).collect()
+            );
+        }
+    };
+    result.into()
+}

pytests/common/mod.rs

@@ -1,4 +1,4 @@
-use std::{future::Future, thread, time::Duration};
+use std::{thread, time::Duration};
 
 use pyo3::prelude::*;
 
@@ -12,25 +12,20 @@ async def sleep_for_1s(sleep_for):
     await sleep_for(1)
 "#;
 
-pub(super) fn test_into_future(
-    py: Python,
-) -> PyResult<impl Future<Output = PyResult<()>> + Send + 'static> {
-    let test_mod: PyObject =
-        PyModule::from_code(py, TEST_MOD, "test_rust_coroutine/test_mod.py", "test_mod")?.into();
-
-    Ok(async move {
-        Python::with_gil(|py| {
-            pyo3_asyncio::into_future(
-                test_mod
-                    .call_method1(py, "py_sleep", (1.into_py(py),))?
-                    .as_ref(py),
-            )
-        })?
-        .await?;
-        Ok(())
-    })
+pub(super) async fn test_into_future() -> PyResult<()> {
+    let fut = Python::with_gil(|py| {
+        let test_mod =
+            PyModule::from_code(py, TEST_MOD, "test_rust_coroutine/test_mod.py", "test_mod")?;
+
+        pyo3_asyncio::into_future(test_mod.call_method1("py_sleep", (1.into_py(py),))?)
+    })?;
+
+    fut.await?;
+
+    Ok(())
 }
 
-pub(super) fn test_blocking_sleep() {
+pub(super) fn test_blocking_sleep() -> PyResult<()> {
     thread::sleep(Duration::from_secs(1));
+    Ok(())
 }

pytests/test_async_std_asyncio.rs

@@ -1,15 +1,10 @@
 mod common;
 
-use std::{future::Future, time::Duration};
+use std::time::Duration;
 
 use async_std::task;
 use pyo3::{prelude::*, wrap_pyfunction};
 
-use pyo3_asyncio::{
-    async_std::testing::{new_sync_test, test_main},
-    testing::Test,
-};
-
 #[pyfunction]
 fn sleep_for(py: Python, secs: &PyAny) -> PyResult<PyObject> {
     let secs = secs.extract()?;
@@ -20,91 +15,53 @@ fn sleep_for(py: Python, secs: &PyAny) -> PyResult<PyObject> {
     })
 }
 
-fn test_into_coroutine(
-    py: Python,
-) -> PyResult<impl Future<Output = PyResult<()>> + Send + 'static> {
-    let sleeper_mod: Py<PyModule> = PyModule::new(py, "rust_sleeper")?.into();
-
-    sleeper_mod
-        .as_ref(py)
-        .add_wrapped(wrap_pyfunction!(sleep_for))?;
-
-    let test_mod: PyObject = PyModule::from_code(
-        py,
-        common::TEST_MOD,
-        "test_rust_coroutine/test_mod.py",
-        "test_mod",
-    )?
-    .into();
-
-    Ok(async move {
-        Python::with_gil(|py| {
-            pyo3_asyncio::into_future(
-                test_mod
-                    .call_method1(py, "sleep_for_1s", (sleeper_mod.getattr(py, "sleep_for")?,))?
-                    .as_ref(py),
-            )
-        })?
-        .await?;
-        Ok(())
-    })
+#[pyo3_asyncio::async_std::test]
+async fn test_into_coroutine() -> PyResult<()> {
+    let fut = Python::with_gil(|py| {
+        let sleeper_mod = PyModule::new(py, "rust_sleeper")?;
+
+        sleeper_mod.add_wrapped(wrap_pyfunction!(sleep_for))?;
+
+        let test_mod = PyModule::from_code(
+            py,
+            common::TEST_MOD,
+            "test_rust_coroutine/test_mod.py",
+            "test_mod",
+        )?;
+
+        pyo3_asyncio::into_future(
+            test_mod.call_method1("sleep_for_1s", (sleeper_mod.getattr("sleep_for")?,))?,
+        )
+    })?;
+
+    fut.await?;
+
+    Ok(())
 }
 
-fn test_async_sleep<'p>(
-    py: Python<'p>,
-) -> PyResult<impl Future<Output = PyResult<()>> + Send + 'static> {
-    let asyncio = PyObject::from(py.import("asyncio")?);
+#[pyo3_asyncio::async_std::test]
+async fn test_async_sleep() -> PyResult<()> {
+    let asyncio =
+        Python::with_gil(|py| py.import("asyncio").map(|asyncio| PyObject::from(asyncio)))?;
 
-    Ok(async move {
-        task::sleep(Duration::from_secs(1)).await;
+    task::sleep(Duration::from_secs(1)).await;
 
-        Python::with_gil(|py| {
-            pyo3_asyncio::into_future(asyncio.as_ref(py).call_method1("sleep", (1.0,))?)
-        })?
-        .await?;
+    Python::with_gil(|py| {
+        pyo3_asyncio::into_future(asyncio.as_ref(py).call_method1("sleep", (1.0,))?)
+    })?
+    .await?;
 
-        Ok(())
-    })
+    Ok(())
 }
 
-fn main() {
-    test_main(
-        "PyO3 Asyncio Test Suite",
-        vec![
-            Test::new_async(
-                "test_async_sleep".into(),
-                Python::with_gil(|py| {
-                    test_async_sleep(py)
-                        .map_err(|e| {
-                            e.print_and_set_sys_last_vars(py);
-                        })
-                        .unwrap()
-                }),
-            ),
-            new_sync_test("test_blocking_sleep".into(), || {
-                common::test_blocking_sleep();
-                Ok(())
-            }),
-            Test::new_async(
-                "test_into_coroutine".into(),
-                Python::with_gil(|py| {
-                    test_into_coroutine(py)
-                        .map_err(|e| {
-                            e.print_and_set_sys_last_vars(py);
-                        })
-                        .unwrap()
-                }),
-            ),
-            Test::new_async(
-                "test_into_future".into(),
-                Python::with_gil(|py| {
-                    common::test_into_future(py)
-                        .map_err(|e| {
-                            e.print_and_set_sys_last_vars(py);
-                        })
-                        .unwrap()
-                }),
-            ),
-        ],
-    )
+#[pyo3_asyncio::async_std::test]
+fn test_blocking_sleep() -> PyResult<()> {
+    common::test_blocking_sleep()
 }
+
+#[pyo3_asyncio::async_std::test]
+async fn test_into_future() -> PyResult<()> {
+    common::test_into_future().await
+}
+
+pyo3_asyncio::async_std::test_main!("PyO3 Asyncio Test Suite for Async-Std Runtime");

pytests/test_async_std_test_macro.rs

@@ -1,8 +1,5 @@
 mod common;
 mod tokio_asyncio;
 
-fn main() {
-    pyo3_asyncio::tokio::init_current_thread();
-
-    tokio_asyncio::test_main("PyO3 Asyncio Tokio Current-Thread Test Suite");
-}
+// TODO: Fix current thread init
+pyo3_asyncio::tokio::test_main!("PyO3 Asyncio Test Suite for Tokio Current-Thread Runtime");