lance_test_macros/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4extern crate proc_macro;
5
6use proc_macro::TokenStream;
7use proc_macro2::TokenStream as Tokens;
8
9use quote::quote;
10use syn::{parse_macro_input, punctuated::Punctuated, FnArg, ItemFn, ReturnType, Token};
11
12// The tracing initialization
13//
14// Note that there are two guards.  The first is for the chrome layer and the
15// second is for the tracing subscriber.  The tuple order is important as the
16// chrome layer guard must be dropped before the subscriber guard.
17fn expand_tracing_init() -> Tokens {
18    quote! {
19      {
20        let trace_level = std::env::var("LANCE_TRACING");
21        if let Ok(trace_level) = trace_level {
22
23          let level_filter = match trace_level.as_str() {
24            "debug" => ::tracing_subscriber::filter::LevelFilter::DEBUG,
25            "info" => ::tracing_subscriber::filter::LevelFilter::INFO,
26            "warn" => ::tracing_subscriber::filter::LevelFilter::WARN,
27            "error" => ::tracing_subscriber::filter::LevelFilter::ERROR,
28            "trace" => ::tracing_subscriber::filter::LevelFilter::TRACE,
29            _ => panic!("Unexpected trace level {}", trace_level),
30          };
31
32          let (chrome_layer, _guard) = tracing_chrome::ChromeLayerBuilder::new().trace_style(tracing_chrome::TraceStyle::Async).include_args(true).build();
33          let subscriber = ::tracing_subscriber::registry::Registry::default();
34          let chrome_layer = ::tracing_subscriber::Layer::with_filter(chrome_layer, level_filter);
35          let subscriber = tracing_subscriber::prelude::__tracing_subscriber_SubscriberExt::with(
36            subscriber, chrome_layer);
37          let sub_guard = ::tracing::subscriber::set_default(subscriber);
38          Some((_guard, sub_guard))
39        } else {
40          None
41        }
42      }
43    }
44}
45
46fn extract_args(inputs: &Punctuated<FnArg, Token![,]>) -> Punctuated<Tokens, Token![,]> {
47    inputs
48        .iter()
49        .map(|arg| match arg {
50            FnArg::Receiver(receiver) => {
51                let slf = receiver.self_token;
52                quote! { #slf }
53            }
54            FnArg::Typed(typed) => {
55                let pat = &typed.pat;
56                quote! { #pat }
57            }
58        })
59        .collect()
60}
61
62// This function parses the wrapped object into a function (tests are functions) and
63// then creates a new wrapped function
64fn expand_wrapper(wrapped_attr: Tokens, wrappee: ItemFn) -> Tokens {
65    let attrs = &wrappee.attrs;
66    let async_ = &wrappee.sig.asyncness;
67    let await_ = if async_.is_some() {
68        quote! {.await}
69    } else {
70        quote! {}
71    };
72    let body = &wrappee.block;
73    let test_name = &wrappee.sig.ident;
74    let inputs = &wrappee.sig.inputs;
75    let args = extract_args(inputs);
76
77    // Note that Rust does not allow us to have a test function with
78    // #[should_panic] that has a non-unit return value.
79    let ret = match &wrappee.sig.output {
80        ReturnType::Default => quote! {},
81        ReturnType::Type(_, type_) => quote! {-> #type_},
82    };
83
84    let tracing_init = expand_tracing_init();
85
86    // Creates a test-scoped init function and then calls the underlying test
87    let result = quote! {
88      #[#wrapped_attr]
89      #(#attrs)*
90      #async_ fn #test_name(#inputs) #ret {
91        #async_ fn test_impl(#inputs) #ret {
92          #body
93        }
94
95        mod init {
96          pub fn init() -> Option<(tracing_chrome::FlushGuard, tracing::subscriber::DefaultGuard)> {
97            #tracing_init
98          }
99        }
100
101        let _guard = init::init();
102        test_impl(#args)#await_
103      }
104    };
105    result
106}
107
108// Note: this is a fork of https://crates.io/crates/test-log
109//
110// The original crate could only configure logging tracing to stdout and this
111// is a good entrypoint for any other lance-specific test behaviors we want to
112// add in the future.
113/// This attribute wraps any existing test attribute (e.g. tokio::test or test)
114/// to configure tracing
115///
116/// Example:
117///
118/// ```rust,ignore
119/// #[lance_test_macros::test(tokio::test)]
120/// async fn test_something() {
121///  ...
122/// }
123/// ```
124///
125/// By default this wrapper will do nothing.  To then get tracing output, set the
126/// LANCE_TRACING environment variable to your desired level (e.g. "debug").
127///
128/// Example:
129///
130/// ```bash
131/// LANCE_TRACING=debug cargo test dataset::tests::test_create_dataset
132/// ```
133///
134/// A .json file will be generated in the current directory that can be loaded into
135/// chrome://tracing or the perfetto UI.
136///
137/// Note: if multiple tests are wrapped and you enable the environment variable then
138/// you will get a separate .json file for each test that is run.  So generally you
139/// only want to set the environment variable when running a single test at a time.
140#[proc_macro_attribute]
141pub fn test(args: TokenStream, input: TokenStream) -> TokenStream {
142    let input = parse_macro_input!(input as ItemFn);
143
144    expand_wrapper(args.into(), input).into()
145}
lance_test_macros/lib.rs

lance_test_macros/
lib.rs
