sdk/retry: support ending the iteration early

I've found this feature to be very useful in https://pkg.go.dev/gotest.tools/v3/poll#WaitOn

I have encountered a few cases where I wanted that same support, so this commit adds it.
This commit is contained in:
Daniel Nephin 2021-04-23 17:01:05 -04:00
parent 034c5c5f8c
commit a2986ebb9c
2 changed files with 45 additions and 7 deletions

View File

@ -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.

View File

@ -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() {