after_test/
lib.rs

1//! The [`cleanup`] macro iterates over all functions marked as tests
2//! in the input token stream and adds a call to the provided function
3//! passed as an attribute to the macro.
4//!
5//! This macro can be added to the top of your test module and will ensure
6//! that the passed cleanup function is called at each test function's end.
7//!
8//! # Example
9//!
10//! ```rust
11//! use after_test::cleanup;
12//!
13//! #[cleanup(my_clean_up)]
14//! #[cfg(test)]
15//! mod tests {
16//!     fn my_clean_up() {
17//!        println!("cleaning up resources");
18//!     }   
19//!
20//!     #[test]
21//!     fn a_test() {}
22//! }
23//! ```
24//!
25//! ```rust
26//! use after_test::cleanup;
27//!
28//! #[cfg(test)]
29//! #[cleanup(|| {println ! ("this will be called at the end of each test")})]
30//! mod tests {
31//!     #[test]
32//!     fn my_test() {}
33//! }
34//! ```
35
36mod attribute;
37
38use crate::attribute::CleanupFunction;
39use proc_macro::TokenStream;
40use proc_macro_error::{emit_error, proc_macro_error};
41use quote::quote;
42use syn::parse_macro_input;
43
44#[proc_macro_error]
45#[proc_macro_attribute]
46pub fn cleanup(attr: TokenStream, item: TokenStream) -> TokenStream {
47    let module = parse_macro_input!(item as syn::ItemMod);
48
49    // Assert the module is marked with the `#[cfg(test)]` attribute
50    assert_test_mod(&module);
51
52    // Parse the cleanup function identifier
53    let clean_fn = parse_macro_input!(attr as CleanupFunction);
54
55    // Add the cleanup function call where needed
56    let module = clean(module, clean_fn);
57
58    quote!(
59        #module
60    )
61    .into()
62}
63
64/// Asserts the module is marked with the `#[cfg(test)]` attribute
65fn assert_test_mod(module: &syn::ItemMod) {
66    let is_test_mod = module.attrs.iter().any(|attr| {
67        attr.meta.path().is_ident("cfg")
68            && attr
69                .meta
70                .require_list()
71                .map(|l| l.tokens.to_string() == "test")
72                .unwrap_or_default()
73    });
74
75    if !is_test_mod {
76        emit_error!(
77            module,
78            "module should be marked with the `#[cfg(test)]` attribute";
79            help = "add `#[cfg(test)]` to the module"
80        );
81    }
82}
83
84/// Adds the cleanup function call on each function marked
85/// with `#[test]` attribute
86fn clean(mut module: syn::ItemMod, clean_fn: CleanupFunction) -> syn::ItemMod {
87    let is_test = |f: &syn::ItemFn| f.attrs.iter().any(|attr| attr.path().is_ident("test"));
88
89    module.content = module.content.map(|(brace, items)| {
90        let n_items = items
91            .into_iter()
92            .filter_map(|i| {
93                let f = match &i {
94                    syn::Item::Fn(f) if is_test(f) => {
95                        let attr = &f.attrs;
96                        let block = &f.block;
97                        let sig = &f.sig;
98                        let vis = &f.vis;
99                        syn::parse2(quote!(
100                            #(#attr)*
101                            #vis #sig {
102                                #block
103                                #clean_fn;
104                            }
105                        ))
106                    }
107                    i => syn::parse2(quote!(#i)),
108                };
109                f.inspect_err(|err| emit_error!(i, err.to_string())).ok()
110            })
111            .collect::<Vec<_>>();
112
113        (brace, n_items)
114    });
115
116    module
117}
after_test/lib.rs

after_test/
lib.rs
