From 7b0a6c0dfaf73bd92f0d3cff8e818e63e2645f50 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 3 Apr 2023 18:16:08 -0700
Subject: [PATCH] Add an 'on_failure' attribute to gpui tests

This lets us perform a finalization step when a randomized test fails.
---
 crates/gpui/src/test.rs               |  2 ++
 crates/gpui_macros/src/gpui_macros.rs | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/crates/gpui/src/test.rs b/crates/gpui/src/test.rs
index d784d43ece..0bf3c333b3 100644
--- a/crates/gpui/src/test.rs
+++ b/crates/gpui/src/test.rs
@@ -45,6 +45,7 @@ pub fn run_test(
         Arc<executor::Deterministic>,
         u64,
     )),
+    on_fail_fn: Option<fn()>,
     fn_name: String,
 ) {
     // let _profiler = dhat::Profiler::new_heap();
@@ -177,6 +178,7 @@ pub fn run_test(
                     if is_randomized {
                         eprintln!("failing seed: {}", atomic_seed.load(SeqCst));
                     }
+                    on_fail_fn.map(|f| f());
                     panic::resume_unwind(error);
                 }
             }
diff --git a/crates/gpui_macros/src/gpui_macros.rs b/crates/gpui_macros/src/gpui_macros.rs
index cabae1ac0a..42cdb66ee3 100644
--- a/crates/gpui_macros/src/gpui_macros.rs
+++ b/crates/gpui_macros/src/gpui_macros.rs
@@ -1,4 +1,5 @@
 use proc_macro::TokenStream;
+use proc_macro2::Ident;
 use quote::{format_ident, quote};
 use std::mem;
 use syn::{
@@ -15,6 +16,7 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
     let mut num_iterations = 1;
     let mut starting_seed = 0;
     let mut detect_nondeterminism = false;
+    let mut on_failure_fn_name = quote!(None);
 
     for arg in args {
         match arg {
@@ -33,6 +35,20 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
                         Some("retries") => max_retries = parse_int(&meta.lit)?,
                         Some("iterations") => num_iterations = parse_int(&meta.lit)?,
                         Some("seed") => starting_seed = parse_int(&meta.lit)?,
+                        Some("on_failure") => {
+                            if let Lit::Str(name) = meta.lit {
+                                let ident = Ident::new(&name.value(), name.span());
+                                on_failure_fn_name = quote!(Some(#ident));
+                            } else {
+                                return Err(TokenStream::from(
+                                    syn::Error::new(
+                                        meta.lit.span(),
+                                        "on_failure argument must be a string",
+                                    )
+                                    .into_compile_error(),
+                                ));
+                            }
+                        }
                         _ => {
                             return Err(TokenStream::from(
                                 syn::Error::new(meta.path.span(), "invalid argument")
@@ -152,6 +168,7 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
                         cx.foreground().run(#inner_fn_name(#inner_fn_args));
                         #cx_teardowns
                     },
+                    #on_failure_fn_name,
                     stringify!(#outer_fn_name).to_string(),
                 );
             }
@@ -187,6 +204,7 @@ pub fn test(args: TokenStream, function: TokenStream) -> TokenStream {
                     #max_retries,
                     #detect_nondeterminism,
                     &mut |cx, _, _, seed| #inner_fn_name(#inner_fn_args),
+                    #on_failure_fn_name,
                     stringify!(#outer_fn_name).to_string(),
                 );
             }