diff --git a/Cargo.lock b/Cargo.lock index 7a675454..404c269f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -464,6 +464,16 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "proc_macro" +version = "0.1.0" +dependencies = [ + "proc-macro2", + "quote", + "rand", + "syn", +] + [[package]] name = "quote" version = "1.0.26" @@ -746,6 +756,7 @@ dependencies = [ "lazy_static", "log", "pretty_assertions", + "proc_macro", "rand", "regex", "regex-syntax", diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 6d48e8b9..47e03284 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -69,6 +69,8 @@ version = "0.4.6" features = ["std"] [dev-dependencies] +proc_macro = { path = "src/tests/proc_macro" } + rand = "0.8" tempfile = "3" pretty_assertions = "0.7.2" diff --git a/cli/src/tests/proc_macro/Cargo.toml b/cli/src/tests/proc_macro/Cargo.toml new file mode 100644 index 00000000..a9a2b146 --- /dev/null +++ b/cli/src/tests/proc_macro/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "proc_macro" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = "1" +quote = "1" +rand = "0.8.5" +syn = { version = "1", features = ["full"] } diff --git a/cli/src/tests/proc_macro/src/lib.rs b/cli/src/tests/proc_macro/src/lib.rs new file mode 100644 index 00000000..d831a75b --- /dev/null +++ b/cli/src/tests/proc_macro/src/lib.rs @@ -0,0 +1,137 @@ +use proc_macro::TokenStream; +use proc_macro2::Span; +use quote::quote; +use syn::{ + parse::{Parse, ParseStream}, + parse_macro_input, Error, Expr, Ident, ItemFn, LitInt, Token, +}; + +#[proc_macro_attribute] +pub fn retry(args: TokenStream, input: TokenStream) -> TokenStream { + let count = parse_macro_input!(args as LitInt); + let input = parse_macro_input!(input as ItemFn); + let attrs = input.attrs.clone(); + let name = input.sig.ident.clone(); + + TokenStream::from(quote! { + #(#attrs),* + fn #name() { + #input + + for i in 0..=#count { + let result = std::panic::catch_unwind(|| { + #name(); + }); + + if result.is_ok() { + return; + } + + if i == #count { + std::panic::resume_unwind(result.unwrap_err()); + } + } + } + }) +} + +#[proc_macro_attribute] +pub fn test_with_seed(args: TokenStream, input: TokenStream) -> TokenStream { + struct Args { + retry: LitInt, + seed: Expr, + seed_fn: Option, + } + + impl Parse for Args { + fn parse(input: ParseStream) -> syn::Result { + let mut retry = None; + let mut seed = None; + let mut seed_fn = None; + + while !input.is_empty() { + let name = input.parse::()?; + match name.to_string().as_str() { + "retry" => { + input.parse::()?; + retry.replace(input.parse()?); + } + "seed" => { + input.parse::()?; + seed.replace(input.parse()?); + } + "seed_fn" => { + input.parse::()?; + seed_fn.replace(input.parse()?); + } + x => { + return Err(Error::new( + name.span(), + format!("Unsupported parameter `{x}`"), + )) + } + } + + if !input.is_empty() { + input.parse::()?; + } + } + + if retry.is_none() { + retry.replace(LitInt::new("0", Span::mixed_site())); + } + + Ok(Args { + retry: retry.expect("`retry` parameter is requred"), + seed: seed.expect("`initial_seed` parameter is required"), + seed_fn, + }) + } + } + + let Args { + retry, + seed, + seed_fn, + } = parse_macro_input!(args as Args); + + let seed_fn = seed_fn.iter(); + + let func = parse_macro_input!(input as ItemFn); + let attrs = func.attrs.clone(); + let name = func.sig.ident.clone(); + + // dbg!(quote::ToTokens::into_token_stream(&func)); + + TokenStream::from(quote! { + #[test] + #(#attrs),* + fn #name() { + #func + + let mut seed = #seed; + + for i in 0..=#retry { + let result = std::panic::catch_unwind(|| { + #name(seed); + }); + + if result.is_ok() { + return; + } + + if i == #retry { + std::panic::resume_unwind(result.unwrap_err()); + } + + #( + seed = #seed_fn(); + )* + + if i < #retry { + println!("\nRetry {}/{} with a new seed {}", i + 1, #retry, seed); + } + } + } + }) +}