Allow applying serial to mod as well

serial_test_derive/Cargo.toml

@@ -19,6 +19,7 @@ proc-macro2 = "1.0.60" # Because of https://github.com/dtolnay/proc-macro2/issue
 
 [dev-dependencies]
 env_logger = "0.10"
+prettyplease = "0.2"
 
 [features]
 async = []

serial_test_derive/src/lib.rs

@@ -10,6 +10,7 @@ use proc_macro::TokenStream;
 use proc_macro2::{Literal, TokenTree};
 use quote::{format_ident, quote, ToTokens, TokenStreamExt};
 use std::ops::Deref;
+use syn::Result as SynResult;
 
 /// Allows for the creation of serialised Rust tests
 /// ````
@@ -193,7 +194,7 @@ pub fn file_parallel(attr: TokenStream, input: TokenStream) -> TokenStream {
 }
 
 // Based off of https://github.com/dtolnay/quote/issues/20#issuecomment-437341743
-#[derive(Default, Debug)]
+#[derive(Default, Debug, Clone)]
 struct QuoteOption<T>(Option<T>);
 
 impl<T: ToTokens> ToTokens for QuoteOption<T> {
@@ -320,11 +321,69 @@ fn fs_parallel_core(
 
 fn core_setup(
     input: proc_macro2::TokenStream,
-    config: Config,
+    config: &Config,
+    prefix: &str,
+    kind: &str,
+) -> proc_macro2::TokenStream {
+    let fn_ast: SynResult<syn::ItemFn> = syn::parse2(input.clone());
+    match fn_ast {
+        Ok(ast) => {
+            return fn_setup(ast, config, prefix, kind);
+        }
+        Err(_) => {
+            // Assume non-fn, skip
+        }
+    };
+    let mod_ast: SynResult<syn::ItemMod> = syn::parse2(input);
+    match mod_ast {
+        Ok(mut ast) => {
+            let new_content = ast.content.clone().and_then(|(brace, items)| {
+                let new_items = items
+                    .into_iter()
+                    .map(|item| match item {
+                        syn::Item::Fn(item_fn)
+                            if item_fn
+                                .attrs
+                                .iter()
+                                .find(|attr| {
+                                    attr.meta
+                                        .path()
+                                        .segments
+                                        .first()
+                                        .unwrap()
+                                        .ident
+                                        .to_string()
+                                        .contains("test")
+                                })
+                                .is_some() =>
+                        {
+                            syn::parse2(fn_setup(item_fn, config, prefix, kind)).unwrap()
+                        }
+                        other => other,
+                    })
+                    .collect();
+                Some((brace, new_items))
+            });
+            if let Some(nc) = new_content {
+                ast.content.replace(nc);
+            }
+            ast.attrs.retain(|attr| {
+                attr.meta.path().segments.first().unwrap().ident.to_string() != "serial"
+            });
+            ast.into_token_stream().into()
+        }
+        Err(_) => {
+            panic!("Attribute applied to something other than mod or fn!");
+        }
+    }
+}
+
+fn fn_setup(
+    ast: syn::ItemFn,
+    config: &Config,
     prefix: &str,
     kind: &str,
 ) -> proc_macro2::TokenStream {
-    let ast: syn::ItemFn = syn::parse2(input).unwrap();
     let asyncness = ast.sig.asyncness;
     if asyncness.is_some() && cfg!(not(feature = "async")) {
         panic!("async testing attempted with async feature disabled in serial_test!");
@@ -337,8 +396,8 @@ fn core_setup(
     };
     let block = ast.block;
     let attrs: Vec<syn::Attribute> = ast.attrs.into_iter().collect();
-    let names = config.names;
-    let path = config.path;
+    let names = config.names.clone();
+    let path = config.path.clone();
     if let Some(ret) = return_type {
         match asyncness {
             Some(_) => {
@@ -401,23 +460,40 @@ fn serial_setup(
     config: Config,
     prefix: &str,
 ) -> proc_macro2::TokenStream {
-    core_setup(input, config, prefix, "serial")
+    core_setup(input, &config, prefix, "serial")
 }
 
 fn parallel_setup(
     input: proc_macro2::TokenStream,
     config: Config,
     prefix: &str,
 ) -> proc_macro2::TokenStream {
-    core_setup(input, config, prefix, "parallel")
+    core_setup(input, &config, prefix, "parallel")
 }
 
 #[cfg(test)]
 mod tests {
     use super::{fs_serial_core, local_serial_core};
+    use proc_macro2::TokenStream;
     use quote::quote;
     use std::iter::FromIterator;
 
+    fn unparse(input: TokenStream) -> String {
+        let item = syn::parse2(input).unwrap();
+        let file = syn::File {
+            attrs: vec![],
+            items: vec![item],
+            shebang: None,
+        };
+
+        prettyplease::unparse(&file)
+    }
+
+    fn compare_streams(first: TokenStream, second: TokenStream) {
+        let f = unparse(first);
+        assert_eq!(f, unparse(second));
+    }
+
     #[test]
     fn test_serial() {
         let attrs = proc_macro2::TokenStream::new();
@@ -432,7 +508,7 @@ mod tests {
                 serial_test::local_serial_core(vec![""], ::std::option::Option::None, || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
     }
 
     #[test]
@@ -449,7 +525,7 @@ mod tests {
                 serial_test::local_serial_core(vec![""], ::std::option::Option::None, || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
     }
 
     #[test]
@@ -470,10 +546,10 @@ mod tests {
             #[should_panic(expected = "Testing panic")]
             #[something_else]
             fn foo () {
-                serial_test::local_serial_core(vec![""], ::std::option::Option::None,  || {} );
+                serial_test::local_serial_core(vec![""], ::std::option::Option::None, || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
     }
 
     #[test]
@@ -527,7 +603,7 @@ mod tests {
                 serial_test::fs_serial_core(vec!["foo"], ::std::option::Option::None, || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
     }
 
     #[test]
@@ -547,7 +623,7 @@ mod tests {
                 serial_test::fs_serial_core(vec![""], ::std::option::Option::None, || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
     }
 
     #[test]
@@ -567,7 +643,7 @@ mod tests {
                 serial_test::fs_serial_core(vec!["foo"], ::std::option::Option::Some("bar_path"), || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
     }
 
     #[test]
@@ -587,7 +663,7 @@ mod tests {
                 serial_test::local_serial_core(vec!["one"], ::std::option::Option::None, || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
     }
 
     #[test]
@@ -607,6 +683,41 @@ mod tests {
                 serial_test::local_serial_core(vec!["one", "two"], ::std::option::Option::None, || {} );
             }
         };
-        assert_eq!(format!("{}", compare), format!("{}", stream));
+        compare_streams(compare, stream);
+    }
+
+    #[test]
+    fn test_mod() {
+        let attrs = proc_macro2::TokenStream::new();
+        let input = quote! {
+            #[cfg(test)]
+            #[serial]
+            mod serial_attr_tests {
+                pub fn foo() {
+                    println!("Nothing");
+                }
+
+                #[test]
+                fn bar() {}
+            }
+        };
+        let stream = local_serial_core(
+            proc_macro2::TokenStream::from_iter(attrs.into_iter()),
+            input,
+        );
+        let compare = quote! {
+            #[cfg(test)]
+            mod serial_attr_tests {
+                pub fn foo() {
+                    println!("Nothing");
+                }
+
+                #[test]
+                fn bar() {
+                    serial_test::local_serial_core(vec![""], ::std::option::Option::None, || {} );
+                }
+            }
+        };
+        compare_streams(compare, stream);
     }
 }

serial_test_test/src/lib.rs

@@ -39,6 +39,8 @@
 //! ```
 
 use lazy_static::lazy_static;
+#[cfg(test)]
+use serial_test::{parallel, serial};
 use std::{
     convert::TryInto,
     env, fs,
@@ -82,6 +84,14 @@ pub fn fs_test_fn(count: usize) {
     assert_eq!(loaded, count);
 }
 
+#[cfg(test)]
+#[serial]
+mod serial_attr_tests {}
+
+#[cfg(test)]
+#[parallel]
+mod parallel_attr_tests {}
+
 #[cfg(test)]
 mod tests {
     use super::{init, test_fn};