test.rs

  1use proc_macro::TokenStream;
  2use proc_macro2::Ident;
  3use quote::{format_ident, quote};
  4use std::mem;
  5use syn::{
  6    AttributeArgs, FnArg, ItemFn, Lit, Meta, MetaList, NestedMeta, PathSegment, Type, parse_quote,
  7    spanned::Spanned,
  8};
  9
 10pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
 11    let args = syn::parse_macro_input!(args as AttributeArgs);
 12    try_test(args, function).unwrap_or_else(|err| err)
 13}
 14
 15fn try_test(args: Vec<NestedMeta>, function: TokenStream) -> Result<TokenStream, TokenStream> {
 16    let mut seeds = Vec::<u64>::new();
 17    let mut max_retries = 0;
 18    let mut num_iterations = 1;
 19    let mut on_failure_fn_name = quote!(None);
 20
 21    for arg in args {
 22        let NestedMeta::Meta(arg) = arg else {
 23            return Err(error_with_message("unexpected literal", arg));
 24        };
 25
 26        let ident = {
 27            let meta_path = match &arg {
 28                Meta::NameValue(meta) => &meta.path,
 29                Meta::List(list) => &list.path,
 30                Meta::Path(path) => return Err(error_with_message("invalid path argument", path)),
 31            };
 32            let Some(ident) = meta_path.get_ident() else {
 33                return Err(error_with_message("unexpected path", meta_path));
 34            };
 35            ident.to_string()
 36        };
 37
 38        match (&arg, ident.as_str()) {
 39            (Meta::NameValue(meta), "retries") => max_retries = parse_usize(&meta.lit)?,
 40            (Meta::NameValue(meta), "iterations") => num_iterations = parse_usize(&meta.lit)?,
 41            (Meta::NameValue(meta), "on_failure") => {
 42                let Lit::Str(name) = &meta.lit else {
 43                    return Err(error_with_message(
 44                        "on_failure argument must be a string",
 45                        &meta.lit,
 46                    ));
 47                };
 48                let segments = name
 49                    .value()
 50                    .split("::")
 51                    .map(|part| PathSegment::from(Ident::new(part, name.span())))
 52                    .collect();
 53                let path = syn::Path {
 54                    leading_colon: None,
 55                    segments,
 56                };
 57                on_failure_fn_name = quote!(Some(#path));
 58            }
 59            (Meta::NameValue(meta), "seed") => seeds = vec![parse_usize(&meta.lit)? as u64],
 60            (Meta::List(list), "seeds") => seeds = parse_u64_array(&list)?,
 61            (Meta::Path(path), _) => {
 62                return Err(error_with_message("invalid path argument", path));
 63            }
 64            (_, _) => {
 65                return Err(error_with_message("invalid argument name", arg));
 66            }
 67        }
 68    }
 69    let seeds = quote!( #(#seeds),* );
 70
 71    let mut inner_fn = syn::parse::<ItemFn>(function).map_err(error_to_stream)?;
 72    let inner_fn_attributes = mem::take(&mut inner_fn.attrs);
 73    let inner_fn_name = format_ident!("_{}", inner_fn.sig.ident);
 74    let outer_fn_name = mem::replace(&mut inner_fn.sig.ident, inner_fn_name.clone());
 75
 76    let mut outer_fn: ItemFn = if inner_fn.sig.asyncness.is_some() {
 77        // Pass to the test function the number of app contexts that it needs,
 78        // based on its parameter list.
 79        let mut cx_vars = proc_macro2::TokenStream::new();
 80        let mut cx_teardowns = proc_macro2::TokenStream::new();
 81        let mut inner_fn_args = proc_macro2::TokenStream::new();
 82        for (ix, arg) in inner_fn.sig.inputs.iter().enumerate() {
 83            if let FnArg::Typed(arg) = arg {
 84                if let Type::Path(ty) = &*arg.ty {
 85                    let last_segment = ty.path.segments.last();
 86                    match last_segment.map(|s| s.ident.to_string()).as_deref() {
 87                        Some("StdRng") => {
 88                            inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(_seed),));
 89                            continue;
 90                        }
 91                        Some("BackgroundExecutor") => {
 92                            inner_fn_args.extend(quote!(gpui::BackgroundExecutor::new(
 93                                std::sync::Arc::new(dispatcher.clone()),
 94                            ),));
 95                            continue;
 96                        }
 97                        _ => {}
 98                    }
 99                } else if let Type::Reference(ty) = &*arg.ty {
100                    if let Type::Path(ty) = &*ty.elem {
101                        let last_segment = ty.path.segments.last();
102                        if let Some("TestAppContext") =
103                            last_segment.map(|s| s.ident.to_string()).as_deref()
104                        {
105                            let cx_varname = format_ident!("cx_{}", ix);
106                            cx_vars.extend(quote!(
107                                let mut #cx_varname = gpui::TestAppContext::build(
108                                    dispatcher.clone(),
109                                    Some(stringify!(#outer_fn_name)),
110                                );
111                            ));
112                            cx_teardowns.extend(quote!(
113                                dispatcher.run_until_parked();
114                                #cx_varname.quit();
115                                dispatcher.run_until_parked();
116                            ));
117                            inner_fn_args.extend(quote!(&mut #cx_varname,));
118                            continue;
119                        }
120                    }
121                }
122            }
123
124            return Err(error_with_message("invalid function signature", arg));
125        }
126
127        parse_quote! {
128            #[test]
129            fn #outer_fn_name() {
130                #inner_fn
131
132                gpui::run_test(
133                    #num_iterations,
134                    &[#seeds],
135                    #max_retries,
136                    &mut |dispatcher, _seed| {
137                        let executor = gpui::BackgroundExecutor::new(std::sync::Arc::new(dispatcher.clone()));
138                        #cx_vars
139                        executor.block_test(#inner_fn_name(#inner_fn_args));
140                        #cx_teardowns
141                    },
142                    #on_failure_fn_name
143                );
144            }
145        }
146    } else {
147        // Pass to the test function the number of app contexts that it needs,
148        // based on its parameter list.
149        let mut cx_vars = proc_macro2::TokenStream::new();
150        let mut cx_teardowns = proc_macro2::TokenStream::new();
151        let mut inner_fn_args = proc_macro2::TokenStream::new();
152        for (ix, arg) in inner_fn.sig.inputs.iter().enumerate() {
153            if let FnArg::Typed(arg) = arg {
154                if let Type::Path(ty) = &*arg.ty {
155                    let last_segment = ty.path.segments.last();
156
157                    if let Some("StdRng") = last_segment.map(|s| s.ident.to_string()).as_deref() {
158                        inner_fn_args.extend(quote!(rand::SeedableRng::seed_from_u64(_seed),));
159                        continue;
160                    }
161                } else if let Type::Reference(ty) = &*arg.ty {
162                    if let Type::Path(ty) = &*ty.elem {
163                        let last_segment = ty.path.segments.last();
164                        match last_segment.map(|s| s.ident.to_string()).as_deref() {
165                            Some("App") => {
166                                let cx_varname = format_ident!("cx_{}", ix);
167                                let cx_varname_lock = format_ident!("cx_{}_lock", ix);
168                                cx_vars.extend(quote!(
169                                    let mut #cx_varname = gpui::TestAppContext::build(
170                                       dispatcher.clone(),
171                                       Some(stringify!(#outer_fn_name))
172                                    );
173                                    let mut #cx_varname_lock = #cx_varname.app.borrow_mut();
174                                ));
175                                inner_fn_args.extend(quote!(&mut #cx_varname_lock,));
176                                cx_teardowns.extend(quote!(
177                                    drop(#cx_varname_lock);
178                                    dispatcher.run_until_parked();
179                                    #cx_varname.update(|cx| { cx.quit() });
180                                    dispatcher.run_until_parked();
181                                ));
182                                continue;
183                            }
184                            Some("TestAppContext") => {
185                                let cx_varname = format_ident!("cx_{}", ix);
186                                cx_vars.extend(quote!(
187                                    let mut #cx_varname = gpui::TestAppContext::build(
188                                        dispatcher.clone(),
189                                        Some(stringify!(#outer_fn_name))
190                                    );
191                                ));
192                                cx_teardowns.extend(quote!(
193                                    dispatcher.run_until_parked();
194                                    #cx_varname.quit();
195                                    dispatcher.run_until_parked();
196                                ));
197                                inner_fn_args.extend(quote!(&mut #cx_varname,));
198                                continue;
199                            }
200                            _ => {}
201                        }
202                    }
203                }
204            }
205
206            return Err(error_with_message("invalid function signature", arg));
207        }
208
209        parse_quote! {
210            #[test]
211            fn #outer_fn_name() {
212                #inner_fn
213
214                gpui::run_test(
215                    #num_iterations,
216                    &[#seeds],
217                    #max_retries,
218                    &mut |dispatcher, _seed| {
219                        #cx_vars
220                        #inner_fn_name(#inner_fn_args);
221                        #cx_teardowns
222                    },
223                    #on_failure_fn_name,
224                );
225            }
226        }
227    };
228    outer_fn.attrs.extend(inner_fn_attributes);
229
230    Ok(TokenStream::from(quote!(#outer_fn)))
231}
232
233fn parse_usize(literal: &Lit) -> Result<usize, TokenStream> {
234    let Lit::Int(int) = &literal else {
235        return Err(error_with_message("expected an usize", literal));
236    };
237    int.base10_parse().map_err(error_to_stream)
238}
239
240fn parse_u64_array(meta_list: &MetaList) -> Result<Vec<u64>, TokenStream> {
241    meta_list
242        .nested
243        .iter()
244        .map(|meta| {
245            if let NestedMeta::Lit(literal) = &meta {
246                parse_usize(literal).map(|value| value as u64)
247            } else {
248                Err(error_with_message("expected an integer", meta.span()))
249            }
250        })
251        .collect()
252}
253
254fn error_with_message(message: &str, spanned: impl Spanned) -> TokenStream {
255    error_to_stream(syn::Error::new(spanned.span(), message))
256}
257
258fn error_to_stream(err: syn::Error) -> TokenStream {
259    TokenStream::from(err.into_compile_error())
260}