feat: support async test functions

.github/workflows/ci.yml

@@ -8,26 +8,21 @@ jobs:
       github.event_name == 'push' ||
       !startsWith(github.event.pull_request.head.label, 'denoland:')
     runs-on: ubuntu-latest
-
     steps:
       - name: Clone repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Install rust
-        uses: hecrj/setup-rust-action@v1.3.4
+        uses: dtolnay/rust-toolchain@stable
         with:
-          rust-version: 1.54.0
-
-      - name: Install clippy and rustfmt
-        run: |
-          rustup component add clippy
-          rustup component add rustfmt
+          toolchain: stable
+          components: clippy,rustfmt
 
       - name: Build
-        run: cargo build --locked --release --all-targets
+        run: cargo build --release --all-targets
 
       - name: Test
-        run: cargo test --locked --release --all-targets
+        run: cargo test --release --all-targets
 
       - name: Lint
         run: cargo clippy --all-targets

.gitignore

@@ -1 +1,2 @@
 target/
+Cargo.lock

Cargo.toml

@@ -2,15 +2,17 @@
 name = "flaky_test"
 version = "0.1.0"
 authors = ["the Deno authors"]
-edition = "2018"
+edition = "2021"
 license = "MIT"
 repository = "https://github.com/denoland/flaky_test"
 description = "atttribute macro for running a flaky test multiple times"
 
-[lib]
-proc-macro = true
+[workspace]
+members = ["impl"]
 
 [dependencies]
-proc-macro2 = "1"
-quote = "1"
-syn = { version = "1", features = ["full"] }
+flaky_test_impl = { path = "impl" }
+futures-util = { version = "0.3", default-features = false, features = ["std"] }
+
+[dev-dependencies]
+tokio = { version = "1", default-features = false, features = ["rt", "rt-multi-thread", "macros"] }

impl/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "flaky_test_impl"
+version = "0.1.0"
+authors = ["the Deno authors"]
+edition = "2021"
+license = "MIT"
+repository = "https://github.com/denoland/flaky_test"
+description = "implementation detail of the `flaky_test` macro"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "1"
+quote = "1"
+syn = { version = "1", features = ["full"] }

impl/src/lib.rs

@@ -0,0 +1,220 @@
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::parse::Parser as _;
+use syn::punctuated::Punctuated;
+use syn::Attribute;
+use syn::ItemFn;
+use syn::Lit;
+use syn::Meta;
+use syn::MetaList;
+use syn::MetaNameValue;
+use syn::NestedMeta;
+use syn::Token;
+
+struct FlakyTestArgs {
+  times: usize,
+  runtime: Runtime,
+}
+
+enum Runtime {
+  Sync,
+  Tokio(Option<Punctuated<NestedMeta, Token![,]>>),
+}
+
+impl Default for FlakyTestArgs {
+  fn default() -> Self {
+    FlakyTestArgs {
+      times: 3,
+      runtime: Runtime::Sync,
+    }
+  }
+}
+
+fn parse_attr(attr: proc_macro2::TokenStream) -> syn::Result<FlakyTestArgs> {
+  let parser = Punctuated::<Meta, Token![,]>::parse_terminated;
+  let punctuated = parser.parse2(attr)?;
+
+  let mut ret = FlakyTestArgs::default();
+
+  for meta in punctuated {
+    match meta {
+      Meta::Path(path) => {
+        if path.is_ident("tokio") {
+          ret.runtime = Runtime::Tokio(None);
+        } else {
+          return Err(syn::Error::new_spanned(path, "expected `tokio`"));
+        }
+      }
+      Meta::NameValue(MetaNameValue {
+        path,
+        lit: Lit::Int(lit_int),
+        ..
+      }) => {
+        if path.is_ident("times") {
+          ret.times = lit_int.base10_parse::<usize>()?;
+        } else {
+          return Err(syn::Error::new_spanned(
+            path,
+            "expected `times = <int>`",
+          ));
+        }
+      }
+      Meta::List(MetaList { path, nested, .. }) => {
+        if path.is_ident("tokio") {
+          ret.runtime = Runtime::Tokio(Some(nested));
+        } else {
+          return Err(syn::Error::new_spanned(path, "expected `tokio`"));
+        }
+      }
+      _ => {
+        return Err(syn::Error::new_spanned(
+          meta,
+          "expected `times = <int>` or `tokio`",
+        ));
+      }
+    }
+  }
+
+  Ok(ret)
+}
+
+/// A flaky test will be run multiple times until it passes.
+///
+/// # Example
+///
+/// ```rust
+/// use flaky_test::flaky_test;
+///
+/// // By default it will be retried up to 3 times.
+/// #[flaky_test]
+/// fn test_default() {
+///  println!("should pass");
+/// }
+///
+/// // The number of max attempts can be adjusted via `times`.
+/// #[flaky_test(times = 5)]
+/// fn usage_with_named_args() {
+///   println!("should pass");
+/// }
+///
+/// # use std::convert::Infallible;
+/// # async fn async_operation() -> Result<i32, Infallible> {
+/// #   Ok(42)
+/// # }
+/// // Async tests can be run by passing `tokio`.
+/// // Make sure `tokio` is added in your `Cargo.toml`.
+/// #[flaky_test(tokio)]
+/// async fn async_test() {
+///   let res = async_operation().await.unwrap();
+///   assert_eq!(res, 42);
+/// }
+///
+/// // `tokio` and `times` can be combined.
+/// #[flaky_test(tokio, times = 5)]
+/// async fn async_test_five_times() {
+///   let res = async_operation().await.unwrap();
+///   assert_eq!(res, 42);
+/// }
+///
+/// // Any arguments that `#[tokio::test]` supports can be specified.
+/// #[flaky_test(tokio(flavor = "multi_thraed", worker_threads = 2))]
+/// async fn async_test_complex() {
+///   let res = async_operation().await.unwrap();
+///   assert_eq!(res, 42);
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn flaky_test(attr: TokenStream, input: TokenStream) -> TokenStream {
+  let attr = proc_macro2::TokenStream::from(attr);
+  let mut input = proc_macro2::TokenStream::from(input);
+
+  match inner(attr, input.clone()) {
+    Err(e) => {
+      input.extend(e.into_compile_error());
+      input.into()
+    }
+    Ok(t) => t.into(),
+  }
+}
+
+fn inner(
+  attr: proc_macro2::TokenStream,
+  input: proc_macro2::TokenStream,
+) -> syn::Result<proc_macro2::TokenStream> {
+  let args = parse_attr(attr)?;
+  let input_fn: ItemFn = syn::parse2(input)?;
+  let attrs = input_fn.attrs.clone();
+
+  match args.runtime {
+    Runtime::Sync => sync(input_fn, attrs, args.times),
+    Runtime::Tokio(tokio_args) => {
+      tokio(input_fn, attrs, args.times, tokio_args)
+    }
+  }
+}
+
+fn sync(
+  input_fn: ItemFn,
+  attrs: Vec<Attribute>,
+  times: usize,
+) -> syn::Result<proc_macro2::TokenStream> {
+  let fn_name = input_fn.sig.ident.clone();
+
+  Ok(quote! {
+    #[test]
+    #(#attrs)*
+    fn #fn_name() {
+      #input_fn
+
+      for i in 0..#times {
+        println!("flaky_test retry {}", i);
+        let r = ::std::panic::catch_unwind(|| {
+          #fn_name();
+        });
+        if r.is_ok() {
+          return;
+        }
+        if i == #times - 1 {
+          ::std::panic::resume_unwind(r.unwrap_err());
+        }
+      }
+    }
+  })
+}
+
+fn tokio(
+  input_fn: ItemFn,
+  attrs: Vec<Attribute>,
+  times: usize,
+  tokio_args: Option<Punctuated<NestedMeta, Token![,]>>,
+) -> syn::Result<proc_macro2::TokenStream> {
+  if input_fn.sig.asyncness.is_none() {
+    return Err(syn::Error::new_spanned(input_fn.sig, "must be `async fn`"));
+  }
+
+  let fn_name = input_fn.sig.ident.clone();
+  let tokio_macro = match tokio_args {
+    Some(args) => quote! { #[::tokio::test(#args)] },
+    None => quote! { #[::tokio::test] },
+  };
+
+  Ok(quote! {
+    #tokio_macro
+    #(#attrs)*
+    async fn #fn_name() {
+      #input_fn
+
+      for i in 0..#times {
+        println!("flaky_test retry {}", i);
+        let fut = ::std::panic::AssertUnwindSafe(#fn_name());
+        let r = <_ as ::flaky_test::futures_util::future::FutureExt>::catch_unwind(fut).await;
+        if r.is_ok() {
+          return;
+        }
+        if i == #times - 1 {
+          ::std::panic::resume_unwind(r.unwrap_err());
+        }
+      }
+    }
+  })
+}