tetcore_test_utils_derive/
lib.rs

1// This file is part of Tetcore.
2
3// Copyright (C) 2020-2021 Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5
6// This program is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// This program is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with this program. If not, see <https://www.gnu.org/licenses/>.
18
19use 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}