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::{AttributeArgs, ItemFn, Meta, NestedMeta, ReturnType, parse_macro_input, parse_quote};
22
23/// Based on <https://github.com/d-e-s-o/test-log>
24/// Copyright (C) 2019-2022 Daniel Mueller <deso@posteo.net>
25/// SPDX-License-Identifier: (Apache-2.0 OR MIT)
26///
27/// Implementation for the `test` macro.
28pub fn test_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
29    let args = parse_macro_input!(attr as AttributeArgs);
30    let input = parse_macro_input!(item as ItemFn);
31
32    let inner_test = match args.as_slice() {
33        [] => parse_quote! { ::core::prelude::v1::test },
34        [NestedMeta::Meta(Meta::Path(path))] => quote! { #path },
35        [NestedMeta::Meta(Meta::List(list))] => quote! { #list },
36        _ => panic!("unsupported attributes supplied: {:?}", args),
37    };
38
39    expand_wrapper(&inner_test, &input)
40}
41
42fn expand_logging_init() -> TokenStream2 {
43    let crate_name = std::env::var("CARGO_PKG_NAME").unwrap();
44    if crate_name == "mz-ore" {
45        quote! {
46          {
47            use crate::test;
48            let _ = test::init_logging();
49          }
50        }
51    } else {
52        quote! {
53          {
54            let _ = ::mz_ore::test::init_logging();
55          }
56        }
57    }
58}
59
60/// Emit code for a wrapper function around a test function.
61fn expand_wrapper(inner_test: &TokenStream2, wrappee: &ItemFn) -> TokenStream {
62    let attrs = &wrappee.attrs;
63    let async_ = &wrappee.sig.asyncness;
64    let await_ = if async_.is_some() {
65        quote! {.await}
66    } else {
67        quote! {}
68    };
69    let body = &wrappee.block;
70    let test_name = &wrappee.sig.ident;
71
72    // Note that Rust does not allow us to have a test function with
73    // #[should_panic] that has a non-unit return value.
74    let ret = match &wrappee.sig.output {
75        ReturnType::Default => quote! {},
76        ReturnType::Type(_, type_) => quote! {-> #type_},
77    };
78
79    let logging_init = expand_logging_init();
80
81    let result = quote! {
82      #[#inner_test]
83      #(#attrs)*
84      #async_ fn #test_name() #ret {
85        #async_ fn test_impl() #ret {
86          #body
87        }
88
89        #logging_init
90
91        test_impl()#await_
92      }
93    };
94    result.into()
95}