1use proc_macro::TokenStream;
19use proc_macro2::TokenStream as TokenStream2;
20use quote::quote;
21use syn::punctuated::Punctuated;
22use syn::{ItemFn, Meta, ReturnType, Token, parse_macro_input, parse_quote};
23
24pub fn test_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
30 let args: Vec<_> =
31 parse_macro_input!(attr with Punctuated::<Meta, Token![,]>::parse_terminated)
32 .into_iter()
33 .collect();
34 let input = parse_macro_input!(item as ItemFn);
35
36 let inner_test = match args.as_slice() {
37 [] => parse_quote! { ::core::prelude::v1::test },
38 [Meta::Path(path)] => quote! { #path },
39 [Meta::List(list)] => quote! { #list },
40 _ => panic!("unsupported attributes supplied: {:?}", args),
41 };
42
43 expand_wrapper(&inner_test, &input)
44}
45
46fn expand_logging_init() -> TokenStream2 {
47 let crate_name = std::env::var("CARGO_PKG_NAME").unwrap();
48 if crate_name == "mz-ore" {
49 quote! {
50 {
51 use crate::test;
52 let _ = test::init_logging();
53 }
54 }
55 } else {
56 quote! {
57 {
58 let _ = ::mz_ore::test::init_logging();
59 }
60 }
61 }
62}
63
64fn expand_wrapper(inner_test: &TokenStream2, wrappee: &ItemFn) -> TokenStream {
66 let attrs = &wrappee.attrs;
67 let async_ = &wrappee.sig.asyncness;
68 let await_ = if async_.is_some() {
69 quote! {.await}
70 } else {
71 quote! {}
72 };
73 let body = &wrappee.block;
74 let test_name = &wrappee.sig.ident;
75
76 let ret = match &wrappee.sig.output {
79 ReturnType::Default => quote! {},
80 ReturnType::Type(_, type_) => quote! {-> #type_},
81 };
82
83 let logging_init = expand_logging_init();
84
85 let result = quote! {
86 #[#inner_test]
87 #(#attrs)*
88 #async_ fn #test_name() #ret {
89 #async_ fn test_impl() #ret {
90 #body
91 }
92
93 #logging_init
94
95 test_impl()#await_
96 }
97 };
98 result.into()
99}