Allow to overwrite init on a per test basis

with `#[test(init=other_init)]`

Author: t-moe

macros/src/attributes/tests.rs

@@ -17,19 +17,19 @@ pub(crate) fn expand(args: TokenStream, input: TokenStream) -> TokenStream {
     let validated_module = validate::ValidatedModule::from_module_and_args(module, macro_args);
 
     let untouched_tokens = &validated_module.untouched_tokens;
-    let init_fn = validated_module.init_func.as_ref().map(|i| &i.func);
     let tests = validated_module
         .tests
         .iter()
         .map(|test| codegen::test(test, &validated_module));
+    let init_fns = validated_module.init_funcs.values().map(|i| &i.func);
 
     let mod_name = format_ident!("{}", validated_module.module_name);
     quote!(
         #[cfg(test)]
         mod #mod_name {
             #(#untouched_tokens)*
 
-            #init_fn
+            #(#init_fns)*
 
             #(#tests)*
         }

macros/src/attributes/tests/codegen/test.rs

@@ -13,13 +13,10 @@ pub(crate) fn test(test: &TestFunc, module: &ValidatedModule) -> TokenStream {
     let mut embassy_task = None;
 
     // Generate the code block that will call init, run the test and check the outcome.
-    let mut test_invocation = call_test_fn(test, module.init_func.as_ref());
+    let init = module.init_function_for_test(test);
+    let mut test_invocation = call_test_fn(test, init);
 
-    let init_is_async = module
-        .init_func
-        .as_ref()
-        .map(|i| i.asyncness)
-        .unwrap_or_default();
+    let init_is_async = init.map(|i| i.asyncness).unwrap_or_default();
 
     // If the test is async or the init function is async, we need to wrap the test invocation in an executor.
     // Result is still a block

macros/src/attributes/tests/parse/function_attributes.rs

@@ -1,11 +1,12 @@
+use darling::FromMeta;
 use proc_macro_error2::abort;
 use syn::spanned::Spanned;
 use syn::{Attribute, ItemFn};
 
 /// Represents the attributes that can be applied to a function in the test module
 pub(crate) enum FuncAttribute {
     Init,
-    Test,
+    Test(TestAttribute),
     ShouldPanic,
     Ignore,
     Timeout(TimeoutAttribute),
@@ -19,7 +20,7 @@ impl FuncAttribute {
         let ident = attr.path().get_ident()?.to_string();
         Some(match ident.as_str() {
             "init" => FuncAttribute::Init,
-            "test" => FuncAttribute::Test,
+            "test" => FuncAttribute::Test(TestAttribute::from_attr(attr)),
             "should_panic" => FuncAttribute::ShouldPanic,
             "ignore" => FuncAttribute::Ignore,
             "timeout" => FuncAttribute::Timeout(TimeoutAttribute::from_attr(attr)),
@@ -56,6 +57,24 @@ impl TimeoutAttribute {
     }
 }
 
+#[derive(Debug, FromMeta, Default)]
+pub(crate) struct TestAttribute {
+    #[darling(default)]
+    pub init: Option<syn::Ident>,
+}
+
+impl TestAttribute {
+    pub fn from_attr(attr: &Attribute) -> Self {
+        match &attr.meta {
+            syn::Meta::Path(_) => TestAttribute::default(),
+            meta => match TestAttribute::from_meta(meta) {
+                Ok(test_attr) => test_attr,
+                Err(e) => abort!(attr, "failed to parse `test` attribute. Must be of the form #[test(init = init_function)]: {}", e),
+            },
+        }
+    }
+}
+
 pub(crate) struct FunctionWithAttributes {
     /// Original function item without the attributes that we recognize
     pub func: ItemFn,

macros/src/attributes/tests/validate/function.rs

@@ -3,6 +3,7 @@ use proc_macro_error2::abort;
 use syn::{Attribute, ItemFn, ReturnType, Type};
 
 pub(crate) struct InitFunc {
+    pub name: String,
     pub func: ItemFn,
     pub state: Option<Type>,
     pub asyncness: bool,
@@ -14,15 +15,15 @@ impl From<FunctionWithAttributes> for InitFunc {
         for (attr, span) in attributes {
             match attr {
                 FuncAttribute::Init => {}
-                FuncAttribute::Test => unreachable!(),
+                FuncAttribute::Test(_) => unreachable!(),
                 _ => abort!(span, "The `#[init]` function can not have this attribute"),
             }
         }
         if check_fn_sig(&func.sig).is_err() || !func.sig.inputs.is_empty() {
             abort!(
-                    func.sig,
-                    "`#[init]` function must have signature `async fn() [-> Type]` (async/return type are optional)",
-                );
+                func.sig,
+                "`#[init]` function must have signature `async fn() [-> Type]` (async/return type are optional)",
+            );
         }
 
         if cfg!(not(feature = "embassy")) && func.sig.asyncness.is_some() {
@@ -37,6 +38,7 @@ impl From<FunctionWithAttributes> for InitFunc {
             ReturnType::Type(.., ty) => Some(*ty.clone()),
         };
         InitFunc {
+            name: func.sig.ident.to_string(),
             asyncness: func.sig.asyncness.is_some(),
             func,
             state,
@@ -52,6 +54,7 @@ pub(crate) struct TestFunc {
     pub ignore: bool,
     pub asyncness: bool,
     pub timeout: Option<u32>,
+    pub custom_init: Option<syn::Ident>,
 }
 
 impl From<FunctionWithAttributes> for TestFunc {
@@ -60,10 +63,11 @@ impl From<FunctionWithAttributes> for TestFunc {
         let mut should_panic = false;
         let mut ignore = false;
         let mut timeout = None;
+        let mut custom_init = None;
         for (attr, _span) in attributes {
             match attr {
                 FuncAttribute::Init => unreachable!(),
-                FuncAttribute::Test => {}
+                FuncAttribute::Test(attr) => custom_init = attr.init,
                 FuncAttribute::ShouldPanic => should_panic = true,
                 FuncAttribute::Ignore => ignore = true,
                 FuncAttribute::Timeout(t) => timeout = Some(t.value),
@@ -99,17 +103,31 @@ impl From<FunctionWithAttributes> for TestFunc {
             should_panic,
             ignore,
             timeout,
+            custom_init,
+        }
+    }
+}
+
+pub(crate) struct OtherFunc(pub FunctionWithAttributes);
+impl From<FunctionWithAttributes> for OtherFunc {
+    fn from(func: FunctionWithAttributes) -> Self {
+        if let Some((_attr, span)) = func.attributes.first() {
+            abort!(
+                span,
+                "Only `#[test]` or `#[init]` functions can have such an attribute"
+            );
         }
+        OtherFunc(func)
     }
 }
 
-pub(crate) enum Func {
+pub(crate) enum AnnotatedFunction {
     Init(InitFunc),
     Test(TestFunc),
-    Other(FunctionWithAttributes),
+    Other(OtherFunc),
 }
 
-impl From<FunctionWithAttributes> for Func {
+impl From<FunctionWithAttributes> for AnnotatedFunction {
     fn from(func: FunctionWithAttributes) -> Self {
         enum FuncKind {
             Init,
@@ -119,8 +137,8 @@ impl From<FunctionWithAttributes> for Func {
         for (attr, span) in &func.attributes {
             match attr {
                 FuncAttribute::Init if func_kind.is_none() => func_kind = Some(FuncKind::Init),
-                FuncAttribute::Test if func_kind.is_none() => func_kind = Some(FuncKind::Test),
-                FuncAttribute::Init | FuncAttribute::Test => {
+                FuncAttribute::Test(_) if func_kind.is_none() => func_kind = Some(FuncKind::Test),
+                FuncAttribute::Init | FuncAttribute::Test(_) => {
                     abort!(
                         span,
                         "A function can only be marked with one of `#[init]` or `#[test]`"
@@ -131,9 +149,9 @@ impl From<FunctionWithAttributes> for Func {
         }
 
         match func_kind {
-            Some(FuncKind::Init) => Func::Init(InitFunc::from(func)),
-            Some(FuncKind::Test) => Func::Test(TestFunc::from(func)),
-            None => Func::Other(func),
+            Some(FuncKind::Init) => AnnotatedFunction::Init(InitFunc::from(func)),
+            Some(FuncKind::Test) => AnnotatedFunction::Test(TestFunc::from(func)),
+            None => AnnotatedFunction::Other(OtherFunc::from(func)),
         }
     }
 }