turbo_crates_testing_proc_macros/
lib.rs

1extern crate proc_macro;
2extern crate syn;
3#[macro_use]
4extern crate quote;
5use proc_macro::TokenStream;
6use syn::parse_macro_input;
7
8#[proc_macro_attribute]
9/// Some tests need extra guard to be executed safely.
10/// It fixes some flakes that are caused by tests that are executed in parallel and are using the same resources.
11///
12/// Make sure to add `#[safe_test]` attribute to each test that needs to be executed safely.
13///
14/// This macro will inject static mutex variable, so that tests are executed safely without impacting each other.
15/// Mutex variable name can be provided as a macro param.
16/// If no param is provided, default name TEST_MUTEX will be used.
17///
18///
19/// Example:
20/// ```rust
21/// #[safe_tests]
22/// pub(crate) mod test {
23///     const TEST_ENV_VAR: &str = "TEST_ENV_VAR";
24///
25///     #[safe_test]
26///     #[test]
27///     fn it_works1() {
28///         std::env::set_var(TEST_ENV_VAR, "test1");
29///         let test_value = std::env::var(TEST_ENV_VAR).unwrap();
30///         assert_eq!(test_value, "test1");
31///     }
32///     #[safe_test]
33///     #[test]
34///     fn it_works2() {
35///         std::env::set_var(TEST_ENV_VAR, "test2");
36///         let test_value = std::env::var(TEST_ENV_VAR).unwrap();
37///         assert_eq!(test_value, "test2");
38///     }
39/// }
40/// ```
41pub fn safe_tests(attr: TokenStream, input: TokenStream) -> TokenStream {
42    let module = parse_macro_input!(input as syn::ItemMod);
43
44    let mutex_name = attr.to_string();
45
46    let mutex_name = if mutex_name.is_empty() {
47        "TEST_MUTEX".to_string()
48    } else {
49        mutex_name
50    };
51    let mutex_var = syn::Ident::new(&mutex_name, proc_macro2::Span::call_site());
52
53    let content = module.content.unwrap().1;
54    proc_macro::TokenStream::from(quote! {
55        // Inject static mutex variable, so that tests are executed safely without impacting each other
56        static #mutex_var: std::sync::Mutex<()> = std::sync::Mutex::new(());
57        #(#content)*
58    })
59    // input
60}
61
62fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
63    tokens.extend(TokenStream::from(error.into_compile_error()));
64    tokens
65}
66
67/// Uses static mutex variable to execute test safely.
68/// Mutex variable name can be provided as a macro param.
69/// If no param is provided, default name TEST_MUTEX will be used.
70#[proc_macro_attribute]
71pub fn safe_test(attr: TokenStream, input: TokenStream) -> TokenStream {
72    let input: syn::ItemFn = match syn::parse(input.clone()) {
73        Ok(it) => it,
74        Err(e) => return token_stream_with_error(input, e),
75    };
76
77    let mutex_name = attr.to_string();
78
79    let mutex_name = if mutex_name.is_empty() {
80        "TEST_MUTEX".to_string()
81    } else {
82        mutex_name
83    };
84    let mutex_var = syn::Ident::new(&mutex_name, proc_macro2::Span::call_site());
85
86    let attrs = &input.attrs;
87    let vis = &input.vis;
88    let sig = &input.sig;
89    let block = &input.block;
90
91    let stmts = &block.stmts;
92
93    let output = quote! {
94        #(#attrs)* #vis #sig {
95            let _guard = #mutex_var.lock().unwrap();
96
97            #(#stmts)*
98        }
99    };
100
101    output.into()
102}