mz_ore_proc/
test.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! test macro with auto-initialized logging
17
18use 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
24/// Based on <https://github.com/d-e-s-o/test-log>
25/// Copyright (C) 2019-2022 Daniel Mueller <deso@posteo.net>
26/// SPDX-License-Identifier: (Apache-2.0 OR MIT)
27///
28/// Implementation for the `test` macro.
29pub 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
64/// Emit code for a wrapper function around a test function.
65fn 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    // Note that Rust does not allow us to have a test function with
77    // #[should_panic] that has a non-unit return value.
78    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}