tetcore_test_utils_derive/
lib.rs1use proc_macro::{Span, TokenStream};
20use proc_macro_crate::crate_name;
21use quote::quote;
22use std::env;
23
24#[proc_macro_attribute]
25pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
26 impl_test(args, item)
27}
28
29fn impl_test(args: TokenStream, item: TokenStream) -> TokenStream {
30 let input = syn::parse_macro_input!(item as syn::ItemFn);
31 let args = syn::parse_macro_input!(args as syn::AttributeArgs);
32
33 parse_knobs(input, args).unwrap_or_else(|e| e.to_compile_error().into())
34}
35
36fn parse_knobs(
37 mut input: syn::ItemFn,
38 args: syn::AttributeArgs,
39) -> Result<TokenStream, syn::Error> {
40 let sig = &mut input.sig;
41 let body = &input.block;
42 let attrs = &input.attrs;
43 let vis = input.vis;
44
45 if sig.inputs.len() != 1 {
46 let msg = "the test function accepts only one argument of type tc_service::TaskExecutor";
47 return Err(syn::Error::new_spanned(&sig, msg));
48 }
49 let (task_executor_name, task_executor_type) = match sig.inputs.pop().map(|x| x.into_value()) {
50 Some(syn::FnArg::Typed(x)) => (x.pat, x.ty),
51 _ => {
52 let msg =
53 "the test function accepts only one argument of type tc_service::TaskExecutor";
54 return Err(syn::Error::new_spanned(&sig, msg));
55 }
56 };
57
58 let crate_name = if env::var("CARGO_PKG_NAME").unwrap() == "tetcore-test-utils" {
59 syn::Ident::new("tetcore_test_utils", Span::call_site().into())
60 } else {
61 let crate_name = crate_name("tetcore-test-utils")
62 .map_err(|e| syn::Error::new_spanned(&sig, e))?;
63
64 syn::Ident::new(&crate_name, Span::call_site().into())
65 };
66
67 let header = {
68 quote! {
69 #[#crate_name::tokio::test(#(#args)*)]
70 }
71 };
72
73 let result = quote! {
74 #header
75 #(#attrs)*
76 #vis #sig {
77 use #crate_name::futures::future::FutureExt;
78
79 let #task_executor_name: #task_executor_type = (|fut, _| {
80 #crate_name::tokio::spawn(fut).map(drop)
81 })
82 .into();
83 let timeout_task = #crate_name::tokio::time::delay_for(
84 std::time::Duration::from_secs(
85 std::env::var("TETCORE_TEST_TIMEOUT")
86 .ok()
87 .and_then(|x| x.parse().ok())
88 .unwrap_or(600))
89 ).fuse();
90 let actual_test_task = async move {
91 #body
92 }
93 .fuse();
94
95 #crate_name::futures::pin_mut!(timeout_task, actual_test_task);
96
97 #crate_name::futures::select! {
98 _ = timeout_task => {
99 panic!("The test took too long!");
100 },
101 _ = actual_test_task => {},
102 }
103 }
104 };
105
106 Ok(result.into())
107}