From 85112bd4d831b744d2e516c1962227aa69c3f6f3 Mon Sep 17 00:00:00 2001
From: Baha Aiman <bahaaiman@google.com>
Date: Tue, 26 Nov 2024 11:27:11 -0800
Subject: [PATCH] test(internal): Add fatal to stop retrying (#11148)

---
 internal/testutil/retry.go | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/internal/testutil/retry.go b/internal/testutil/retry.go
index 7308d7364cfb..3df140794fbd 100644
--- a/internal/testutil/retry.go
+++ b/internal/testutil/retry.go
@@ -45,6 +45,12 @@ func Retry(t *testing.T, maxAttempts int, sleep time.Duration, f func(r *R)) boo
 			t.Fail()
 		}
 
+		if r.fatal {
+			t.Logf("Fatal failure after %d attempts:%s", attempt, r.log.String())
+			t.Fail()
+			return false
+		}
+
 		time.Sleep(sleep)
 	}
 	return false
@@ -70,6 +76,11 @@ func RetryWithoutTest(maxAttempts int, sleep time.Duration, f func(r *R)) bool {
 			return false
 		}
 
+		if r.fatal {
+			r.Logf("Fatal failure after %d attempts:%s", attempt, r.log.String())
+			return false
+		}
+
 		time.Sleep(sleep)
 	}
 	return false
@@ -81,6 +92,7 @@ type R struct {
 	Attempt int
 
 	failed bool
+	fatal  bool
 	log    *bytes.Buffer
 }
 
@@ -95,6 +107,12 @@ func (r *R) Errorf(s string, v ...interface{}) {
 	r.Fail()
 }
 
+// Fatalf is equivalent to Errorf but will not retry.
+func (r *R) Fatalf(s string, v ...interface{}) {
+	r.Errorf(s, v...)
+	r.fatal = true
+}
+
 // Logf formats its arguments and records it in the error log.
 // The text is only printed for the final unsuccessful run or the first successful run.
 func (r *R) Logf(s string, v ...interface{}) {