Consolidated async-std and tokio test_main proc macros

Author: Andrew J Westlake

pyo3-asyncio-macros/src/lib.rs

@@ -153,40 +153,6 @@ pub fn async_std_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
     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: 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::async_std::testing::test_main(
-                #suite_name,
-                inventory::iter::<Test>().map(|test| test.clone()).collect()
-            );
-        }
-    };
-    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 {
@@ -256,12 +222,12 @@ impl Parse for Item {
     }
 }
 
-struct TokioTestMainArgs {
+struct TestMainArgs {
     attrs: Vec<Attribute>,
     suite_name: String,
 }
 
-impl Parse for TokioTestMainArgs {
+impl Parse for TestMainArgs {
     fn parse(input: ParseStream<'_>) -> Result<Self> {
         let mut args: syn::punctuated::Punctuated<Item, syn::Token![,]> =
             input.parse_terminated(Item::parse)?;
@@ -283,9 +249,8 @@ impl Parse for TokioTestMainArgs {
 
 #[cfg(not(test))]
 #[proc_macro]
-pub fn tokio_test_main(args: TokenStream) -> TokenStream {
-    let TokioTestMainArgs { attrs, suite_name } =
-        syn::parse_macro_input!(args as TokioTestMainArgs);
+pub fn test_main(args: TokenStream) -> TokenStream {
+    let TestMainArgs { attrs, suite_name } = syn::parse_macro_input!(args as TestMainArgs);
 
     let result = quote! {
         #[derive(Clone)]

pytests/test_async_std_asyncio.rs

@@ -64,4 +64,7 @@ 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");
+pyo3_asyncio::testing::test_main!(
+    #[pyo3_asyncio::async_std::main],
+    "PyO3 Asyncio Test Suite for Async-Std Runtime"
+);

pytests/test_tokio_current_thread_asyncio.rs

@@ -2,7 +2,7 @@ mod common;
 mod tokio_asyncio;
 
 // TODO: Fix current thread init
-pyo3_asyncio::tokio::test_main!(
+pyo3_asyncio::testing::test_main!(
     #[pyo3_asyncio::tokio::main(flavor = "current_thread")],
     "PyO3 Asyncio Tokio Current Thread Test Suite"
 );

pytests/test_tokio_multi_thread_asyncio.rs

@@ -1,7 +1,7 @@
 mod common;
 mod tokio_asyncio;
 
-pyo3_asyncio::tokio::test_main!(
+pyo3_asyncio::testing::test_main!(
     #[pyo3_asyncio::tokio::main],
     "PyO3 Asyncio Tokio Multi Thread Test Suite"
 );

src/async_std.rs

@@ -12,9 +12,6 @@ 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;
 
 impl JoinError for AsyncStdJoinError {

src/testing.rs

@@ -17,6 +17,10 @@ use clap::{App, Arg};
 use futures::stream::{self, StreamExt};
 use pyo3::prelude::*;
 
+/// <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>
+#[cfg(feature = "attributes")]
+pub use pyo3_asyncio_macros::test_main;
+
 /// Args that should be provided to the test program
 ///
 /// These args are meant to mirror the default test harness's args.

src/tokio.rs

@@ -19,9 +19,6 @@ pub use pyo3_asyncio_macros::tokio_main as main;
 /// <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 tokio runtime and runs an async fn as main
 #[cfg(feature = "attributes")]
 pub use pyo3_asyncio_macros::tokio_test as test;
-/// <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>
-#[cfg(feature = "attributes")]
-pub use pyo3_asyncio_macros::tokio_test_main as test_main;
 
 static TOKIO_RUNTIME: OnceCell<Runtime> = OnceCell::new();