refactor(macros): reuse code

Author: Berrysoft

compio-macros/src/item_fn.rs

@@ -112,6 +112,7 @@ pub(crate) struct RawBodyItemFn {
     pub vis: Visibility,
     pub sig: Signature,
     pub body: TokenStream,
+    pub test: bool,
 }
 
 impl RawBodyItemFn {
@@ -122,14 +123,22 @@ impl RawBodyItemFn {
             vis,
             sig,
             body,
+            test: false,
         }
     }
 
     pub fn set_args(&mut self, args: RawAttr) {
         self.args = args;
     }
 
+    pub fn set_test(&mut self, test: bool) {
+        self.test = test;
+    }
+
     pub fn emit_fn_to_tokens(&self, tokens: &mut TokenStream) {
+        if self.test {
+            tokens.append_all(quote!(#[test]));
+        }
         tokens.append_all(
             self.attrs
                 .iter()

compio-macros/src/lib.rs

@@ -9,8 +9,6 @@ mod item_fn;
 
 mod main_fn;
 
-mod test_fn;
-
 use proc_macro::TokenStream;
 use proc_macro_crate::{FoundCrate, crate_name};
 use proc_macro2::{Ident, Span};
@@ -27,8 +25,9 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 
 #[proc_macro_attribute]
 pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
-    parse_macro_input!(item as test_fn::CompioTest)
+    parse_macro_input!(item as main_fn::CompioMain)
         .with_args(parse_macro_input!(args as item_fn::RawAttr))
+        .with_test(true)
         .into_token_stream()
         .into()
 }

compio-macros/src/main_fn.rs

@@ -11,6 +11,11 @@ impl CompioMain {
         self.0.set_args(args);
         self
     }
+
+    pub fn with_test(mut self, test: bool) -> Self {
+        self.0.set_test(test);
+        self
+    }
 }
 
 impl Parse for CompioMain {