tokio_test_shutdown_timeout/
lib.rs

1use proc_macro2::{Span, TokenStream};
2use syn::{
3    Ident, ItemFn, LitInt,
4};
5use quote::quote;
6
7
8#[proc_macro_attribute]
9pub fn test(
10    attr: proc_macro::TokenStream,
11    item: proc_macro::TokenStream,
12) -> proc_macro::TokenStream {
13
14    let timeout_param: Option<LitInt> = syn::parse_macro_input!(attr);
15
16    let timeout_sec: u64 = if let Some(to) = timeout_param {
17        to.base10_parse().unwrap()
18    } else {
19        5
20    };
21
22    let fn_node = syn::parse_macro_input!(item);
23    test_impl(fn_node, timeout_sec).into()
24}
25
26fn test_impl(mut fn_node: ItemFn, timeout_sec: u64) -> TokenStream {
27    let test_fn_ident = fn_node.sig.ident.clone();
28    let test_fn_ret = fn_node.sig.output.clone();
29    let new_name = format!("__internal_{}", &fn_node.sig.ident);
30    fn_node.sig.ident = Ident::new(&new_name, Span::call_site());
31
32    let mut internal_fn = fn_node;
33    let internal_fn_ident = &internal_fn.sig.ident;
34    let visibility = internal_fn.vis.clone();
35    let attrs = std::mem::take(&mut internal_fn.attrs);
36
37    let res = quote! {
38        #[test]
39        #(#attrs) *
40        #visibility fn #test_fn_ident () #test_fn_ret {
41            #internal_fn
42
43            let __internal_runtime = ::tokio::runtime::Runtime::new().unwrap();
44            let __catch_res = ::std::panic::catch_unwind(|| __internal_runtime.block_on(#internal_fn_ident()));
45            __internal_runtime.shutdown_timeout(::std::time::Duration::from_secs(#timeout_sec));
46            match __catch_res {
47                Ok(__x) => __x,
48                Err(__panic_e) => ::std::panic::resume_unwind(__panic_e),
49            }
50        }
51    };
52
53    res
54}
55
56#[cfg(test)]
57mod unittests {
58    use syn::{parse_quote, ItemFn};
59
60    #[test]
61    fn simple_test() {
62        let testfn: ItemFn = parse_quote!{
63            async fn mytestfn() -> Result<(), ()> {
64                assert_eq!(false, false);
65                Ok(())
66            }
67        };
68
69        _ = super::test_impl(testfn, 5);
70    }
71
72}
tokio_test_shutdown_timeout/lib.rs

tokio_test_shutdown_timeout/
lib.rs
