From a7070201717611c01bb60b3b7f7face6bb13a1f6 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Wed, 24 Feb 2021 14:07:58 +0600
Subject: [PATCH] fix rt_test macros

---
 ntex-macros/Cargo.toml |  2 +-
 ntex-macros/src/lib.rs | 58 ++++++++++++++++++++++++++++++++++++++++++
 ntex/src/lib.rs        |  2 +-
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/ntex-macros/Cargo.toml b/ntex-macros/Cargo.toml
index deca46a3..5a93fe86 100644
--- a/ntex-macros/Cargo.toml
+++ b/ntex-macros/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ntex-macros"
-version = "0.1.2"
+version = "0.1.3"
 description = "ntex proc macros"
 readme = "README.md"
 authors = ["ntex contributors <team@ntex.rs>"]
diff --git a/ntex-macros/src/lib.rs b/ntex-macros/src/lib.rs
index e6e3b3be..8c6c29d2 100644
--- a/ntex-macros/src/lib.rs
+++ b/ntex-macros/src/lib.rs
@@ -279,3 +279,61 @@ pub fn rt_test(_: TokenStream, item: TokenStream) -> TokenStream {
 
     result.into()
 }
+
+/// Marks async test function to be executed by ntex runtime.
+///
+/// ## Usage
+///
+/// ```no_run
+/// #[ntex::test]
+/// async fn my_test() {
+///     assert!(true);
+/// }
+/// ```
+#[doc(hidden)]
+#[proc_macro_attribute]
+pub fn rt_test2(_: TokenStream, item: TokenStream) -> TokenStream {
+    let input = syn::parse_macro_input!(item as syn::ItemFn);
+
+    let ret = &input.sig.output;
+    let name = &input.sig.ident;
+    let body = &input.block;
+    let attrs = &input.attrs;
+    let mut has_test_attr = false;
+
+    for attr in attrs {
+        if attr.path.is_ident("test") {
+            has_test_attr = true;
+        }
+    }
+
+    if input.sig.asyncness.is_none() {
+        return syn::Error::new_spanned(
+            input.sig.fn_token,
+            format!("only async fn is supported, {}", input.sig.ident),
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    let result = if has_test_attr {
+        quote! {
+            #(#attrs)*
+            fn #name() #ret {
+                ntex_rt::System::new("test")
+                    .block_on(async { #body })
+            }
+        }
+    } else {
+        quote! {
+            #[test]
+            #(#attrs)*
+            fn #name() #ret {
+                ntex_rt::System::new("test")
+                    .block_on(async { #body })
+            }
+        }
+    };
+
+    result.into()
+}
diff --git a/ntex/src/lib.rs b/ntex/src/lib.rs
index 37348b9b..b5e78bea 100644
--- a/ntex/src/lib.rs
+++ b/ntex/src/lib.rs
@@ -31,7 +31,7 @@ extern crate derive_more;
 pub use ntex_macros::{rt_main as main, rt_test as test};
 
 #[cfg(test)]
-pub(crate) use ntex_macros::rt_test;
+pub(crate) use ntex_macros::rt_test2 as rt_test;
 
 pub mod channel;
 pub mod connect;