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}