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