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:
parent
f169708bb0
commit
d0b2a4a9d3
|
@ -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.
|
||||
|
|
|
@ -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() {
|
||||
|
|
Loading…
Reference in New Issue