fickle_macros/

lib.rs

use std::fmt::Display;
use std::str::FromStr;

use proc_macro::TokenStream;
use quote::quote;
use syn::punctuated::Punctuated;
use syn::spanned::Spanned;
use syn::{parse_macro_input, Expr, Item, Meta, Token};

#[allow(clippy::test_attr_in_doctest)]
/// Mark a test as fickle (AKA flaky).
///
/// This macro will run your test multiple times if it continues to fail. By default, it will only
/// retry once. To override use the `retries` argument.
///
/// ```
/// #[test]
/// #[fickle]
/// fn my_fickle_test() {}
///
/// #[test]
/// #[fickle(retries=2)]
/// fn more_fickle() {}
/// ```
///
/// There is also the option to set a minimum number of passes. For example, to make sure it passes
/// at least 2 out of every 3 tries, you can use the `passes` argument like so
///
/// ```
/// #[test]
/// #[fickle(retries=2, passes=2)]
/// fn another() {}
/// ```
///
/// It also works with tests that return [`Result`]s.
///
/// ```
/// #[test]
/// #[fickle]
/// fn my_fickle_test() -> Result<(), ()> {
///     Ok(())
/// }
/// ```
#[proc_macro_attribute]
pub fn fickle(attr: TokenStream, item: TokenStream) -> TokenStream {
    // parse the attr args
    let mut n_retries: Option<usize> = None;
    let mut n_passes: Option<usize> = None;
    let attrinp = parse_macro_input!(attr with Punctuated<Meta, Token![,]>::parse_terminated);
    let attrspan = attrinp.span();
    for a in attrinp {
        match a {
            Meta::NameValue(nv) => {
                if nv.path.is_ident("retries") {
                    match get_int(nv.value) {
                        Ok(r) => n_retries = Some(r),
                        Err(err) => return err.into_compile_error().into(),
                    }
                } else if nv.path.is_ident("passes") {
                    match get_int(nv.value) {
                        Ok(r) => n_passes = Some(r),
                        Err(err) => return err.into_compile_error().into(),
                    }
                } else {
                    return syn::Error::new(nv.path.span(), "Unrecognized argument")
                        .into_compile_error()
                        .into();
                }
            }
            _ => {
                return syn::Error::new(a.span(), "Not a name value (`name = value`)")
                    .into_compile_error()
                    .into();
            }
        }
    }

    let fickle_obj = match (n_retries, n_passes) {
        (Some(retries), Some(passes)) => {
            if passes > retries + 1 {
                return syn::Error::new(
                    attrspan,
                    format!(
                        "Requested a minimum of {} passes with a maximum of {} runs",
                        passes,
                        retries + 1
                    ),
                )
                .into_compile_error()
                .into();
            }
            quote! {
                let fickle = fickle::Fickle::new(#retries, #passes).unwrap();
            }
        }
        (Some(retries), None) => quote! {
            let fickle = fickle::Fickle::new_retries(#retries).unwrap();
        },
        (None, Some(passes)) => {
            if passes > 2 {
                return syn::Error::new(
                    attrspan,
                    format!(
                        "Requested a minimum of {} passes with a maximum of {} runs",
                        passes, 2
                    ),
                )
                .into_compile_error()
                .into();
            }
            quote! {
                let fickle = fickle::Fickle::new_passes(#passes).unwrap();
            }
        }
        (None, None) => quote! {
            let fickle = fickle::Fickle::default();
        },
    };

    // parse the func
    let iteminput = parse_macro_input!(item as Item);
    let testfn = match iteminput {
        Item::Fn(func) => func,
        _ => {
            return syn::Error::new(
                iteminput.span(),
                "Expected #[fickle] to annotate a function",
            )
            .into_compile_error()
            .into();
        }
    };
    let attrs = testfn.attrs;
    let vis = testfn.vis;
    let sig = testfn.sig;
    let stmts = testfn.block.stmts;

    let func_body = match &sig.output {
        syn::ReturnType::Default => {
            // these types of tests panic...
            quote! {
                {
                    #fickle_obj
                    let block: fn() -> () = || {
                        #(#stmts)*
                    };
                    fickle.run(block).unwrap();
                }
            }
        }
        syn::ReturnType::Type(_rarrow, typ) => match **typ {
            syn::Type::Path(_) => {
                // this would be where tests that return Results end up
                quote! {
                    {
                        #fickle_obj
                        let block: fn() -> #typ = || {
                            #(#stmts)*
                        };
                        match fickle.run(block) {
                            Ok(res) => Ok(res),
                            Err(fails) => panic!("{}", fails)
                        }
                    }
                }
            }
            _ => {
                return syn::Error::new(typ.span(), "Not able to handle return type")
                    .into_compile_error()
                    .into();
            }
        },
    };

    let expanded = quote! {
        #(#attrs)*
        #vis #sig
        #func_body
    };
    expanded.into()
}

fn get_int<T>(expr: Expr) -> syn::Result<T>
where
    T: FromStr,
    <T as FromStr>::Err: Display,
{
    match expr {
        Expr::Lit(lit) => match lit.lit {
            syn::Lit::Int(litint) => match litint.base10_parse() {
                Ok(r) => Ok(r),
                Err(_) => Err(syn::Error::new(litint.span(), "Could not parse to a usize")),
            },
            _ => Err(syn::Error::new(lit.lit.span(), "Not an integer")),
        },
        _ => Err(syn::Error::new(expr.span(), "Not an literal")),
    }
}