wick_test_logger/
lib.rs

1// Copyright (C) 2019-2021 Daniel Mueller <deso@posteo.net>
2// SPDX-License-Identifier: (Apache-2.0 OR MIT)
3
4extern crate proc_macro;
5
6use proc_macro::TokenStream;
7use proc_macro2::{Ident, Span, TokenStream as Tokens};
8use proc_macro_crate::{crate_name, FoundCrate};
9use quote::quote;
10use syn::parse::Parser;
11use syn::punctuated::Punctuated;
12use syn::{parse_macro_input, parse_quote, ItemFn, Meta, ReturnType, Token};
13
14/// A procedural macro for the `test` attribute.
15///
16/// The attribute can be used to define a test that has the `env_logger`
17/// and/or `tracing` crates initialized (depending on the features used).
18///
19/// # Example
20///
21/// Specify the attribute on a per-test basis:
22/// ```rust
23/// # // doctests seemingly run in a slightly different environment where
24/// # // `super`, which is what our macro makes use of, is not available.
25/// # // By having a fake module here we work around that problem.
26/// # #[cfg(feature = "log")]
27/// # mod fordoctest {
28/// # use logging::info;
29/// # // Note that no test would actually run, regardless of `no_run`,
30/// # // because we do not invoke the function.
31/// #[test_logger::test]
32/// fn it_works() {
33///   info!("Checking whether it still works...");
34///   assert_eq!(2 + 2, 4);
35///   info!("Looks good!");
36/// }
37/// # }
38/// ```
39///
40/// It can be very convenient to convert over all tests by overriding
41/// the `#[test]` attribute on a per-module basis:
42/// ```rust,no_run
43/// # mod fordoctest {
44/// use test_logger::test;
45///
46/// #[test]
47/// fn it_still_works() {
48///   // ...
49/// }
50/// # }
51/// ```
52///
53/// You can also wrap another attribute. For example, suppose you use
54/// [`#[tokio::test]`](https://docs.rs/tokio/1.4.0/tokio/attr.test.html)
55/// to run async tests:
56/// ```
57/// # mod fordoctest {
58/// use test_logger::test;
59///
60/// #[test(tokio::test)]
61/// async fn it_still_works() {
62///   // ...
63/// }
64/// # }
65/// ```
66#[proc_macro_attribute]
67pub fn test(attr: TokenStream, item: TokenStream) -> TokenStream {
68  let parser = Punctuated::<Meta, Token![,]>::parse_terminated;
69  let args = match parser.parse2(attr.into()) {
70    Ok(args) => args,
71    Err(e) => panic!("{}", e),
72  };
73
74  let input = parse_macro_input!(item as ItemFn);
75
76  let args = args.into_iter().collect::<Vec<_>>();
77
78  let inner_test = match args.as_slice() {
79    [] => parse_quote! { ::core::prelude::v1::test },
80    [Meta::Path(path)] => quote! {#path},
81    [Meta::List(list)] => {
82      let path = &list.path;
83      let args = &list.tokens;
84      quote! { #path(#args) }
85    }
86    _ => {
87      panic!("unsupported attributes supplied: {}", quote! { args })
88    }
89  };
90
91  expand_wrapper(inner_test, &input)
92}
93
94/// Expand the initialization code for the `log` crate.
95fn expand_logging_init() -> Tokens {
96  let found_crate = crate_name("wick-logger").expect("wick-logger needs to be added in `Cargo.toml`");
97
98  match found_crate {
99    FoundCrate::Itself => quote! {
100      let logging_options = crate::LoggingOptionsBuilder::default()
101        .app_name("test")
102        .otlp_endpoint(std::env::var("OTLP_ENDPOINT").ok())
103        .levels(crate::LogFilters::with_level(crate::LogLevel::Trace))
104        .build()
105        .unwrap();
106      let __guard = crate::init_test(&logging_options);
107    },
108    FoundCrate::Name(name) => {
109      let ident = Ident::new(&name, Span::call_site());
110
111      quote! {
112        let logging_options = #ident::LoggingOptionsBuilder::default()
113        .app_name("test")
114        .otlp_endpoint(std::env::var("OTLP_ENDPOINT").ok())
115        .levels(#ident::LogFilters::with_level(#ident::LogLevel::Trace))
116        .build()
117        .unwrap();
118        let __guard = #ident::init_test(&logging_options);
119      }
120    }
121  }
122}
123
124/// Emit code for a wrapper function around a test function.
125fn expand_wrapper(inner_test: Tokens, wrappee: &ItemFn) -> TokenStream {
126  let attrs = &wrappee.attrs;
127  let async_ = &wrappee.sig.asyncness;
128  let await_ = if async_.is_some() {
129    quote! {.instrument(span).await}
130  } else {
131    quote! {}
132  };
133  let enter_ = if async_.is_some() {
134    quote! {use tracing::Instrument;}
135  } else {
136    quote! {let _guard = span.enter();}
137  };
138  let exit_ = if async_.is_some() {
139    quote! {
140      // TODO Find a better way to ensure the opentel exporter has flushed.
141      tokio::time::sleep(std::time::Duration::from_millis(200)).await;
142    }
143  } else {
144    quote! {
145      drop(_guard);
146    }
147  };
148  let body = &wrappee.block;
149  let test_name = &wrappee.sig.ident;
150
151  // Note that Rust does not allow us to have a test function with
152  // #[should_panic] that has a non-unit return value.
153  let ret = match &wrappee.sig.output {
154    ReturnType::Default => quote! {},
155    ReturnType::Type(_, type_) => quote! {-> #type_},
156  };
157
158  let logging_init = expand_logging_init();
159
160  let result = quote! {
161    #[#inner_test]
162    #(#attrs)*
163    #async_ fn #test_name() #ret {
164      #async_ fn test_impl() #ret {
165        #body
166      }
167      #logging_init
168      let span = tracing::info_span!(stringify!(#test_name));
169      #enter_
170      let result = test_impl()#await_;
171      if let Err(e) = &result {
172        tracing::error!(error = ?e, "test failed");
173      }
174      #exit_
175      if let Some(guard) = __guard { guard.teardown() } ;
176      result
177    }
178  };
179
180  result.into()
181}
wick_test_logger/lib.rs

wick_test_logger/
lib.rs
