diff --git a/sdk/testutil/retry/retry.go b/sdk/testutil/retry/retry.go index b59bd1c4c..e2b3efc1f 100644 --- a/sdk/testutil/retry/retry.go +++ b/sdk/testutil/retry/retry.go @@ -34,6 +34,7 @@ type Failer interface { // R provides context for the retryer. type R struct { fail bool + done bool output []string } @@ -77,6 +78,12 @@ func (r *R) log(s string) { r.output = append(r.output, decorate(s)) } +// Stop retrying, and fail the test with the specified error. +func (r *R) Stop(err error) { + r.log(err.Error()) + r.done = true +} + func decorate(s string) string { _, file, line, ok := runtime.Caller(3) if ok { @@ -120,6 +127,16 @@ func dedup(a []string) string { func run(r Retryer, t Failer, f func(r *R)) { t.Helper() rr := &R{} + + fail := func() { + t.Helper() + out := dedup(rr.output) + if out != "" { + t.Log(out) + } + t.FailNow() + } + for r.Continue() { func() { defer func() { @@ -129,17 +146,17 @@ func run(r Retryer, t Failer, f func(r *R)) { }() f(rr) }() - if !rr.fail { + + switch { + case rr.done: + fail() + return + case !rr.fail: return } rr.fail = false } - - out := dedup(rr.output) - if out != "" { - t.Log(out) - } - t.FailNow() + fail() } // DefaultFailer provides default retry.Run() behavior for unit tests. diff --git a/sdk/testutil/retry/retry_test.go b/sdk/testutil/retry/retry_test.go index 95186374e..31923a0bf 100644 --- a/sdk/testutil/retry/retry_test.go +++ b/sdk/testutil/retry/retry_test.go @@ -1,6 +1,7 @@ package retry import ( + "fmt" "testing" "time" @@ -52,15 +53,35 @@ func TestRunWith(t *testing.T) { require.Equal(t, 3, iter) require.Equal(t, 1, ft.fails) }) + + t.Run("Stop ends the retrying", func(t *testing.T) { + ft := &fakeT{} + iter := 0 + RunWith(&Counter{Count: 5, Wait: time.Millisecond}, ft, func(r *R) { + iter++ + if iter == 2 { + r.Stop(fmt.Errorf("do not proceed")) + } + r.Fatalf("not yet") + }) + + require.Equal(t, 2, iter) + require.Equal(t, 1, ft.fails) + require.Len(t, ft.out, 1) + require.Contains(t, ft.out[0], "not yet\n") + require.Contains(t, ft.out[0], "do not proceed\n") + }) } type fakeT struct { fails int + out []string } func (f *fakeT) Helper() {} func (f *fakeT) Log(args ...interface{}) { + f.out = append(f.out, fmt.Sprint(args...)) } func (f *fakeT) FailNow() {