bbgo_origin/pkg/util/retry_test.go

107 lines
2.4 KiB
Go
Raw Normal View History

2020-06-08 02:42:50 +00:00
package util
import (
"context"
"fmt"
"strings"
"testing"
"time"
"github.com/pkg/errors"
"github.com/stretchr/testify/assert"
)
func addAndCheck(a *int, target int) error {
if *a++; *a == target {
return nil
} else {
return fmt.Errorf("a is not %v. It is %v\n", target, *a)
}
}
func TestRetry(t *testing.T) {
type test struct {
input int
targetNum int
ans int
ansErr error
}
tests := []test{
{input: 0, targetNum: 3, ans: 3, ansErr: nil},
{input: 0, targetNum: 10, ans: 3, ansErr: errors.New("failed in retry")},
}
for _, tc := range tests {
errHandled := false
err := Retry(context.Background(), 3, 1*time.Second, func() error {
return addAndCheck(&tc.input, tc.targetNum)
}, func(e error) { errHandled = true })
assert.Equal(t, true, errHandled)
if tc.ansErr == nil {
assert.NoError(t, err)
} else {
assert.Contains(t, err.Error(), tc.ansErr.Error())
}
assert.Equal(t, tc.ans, tc.input)
}
}
func TestRetryWithPredicator(t *testing.T) {
type test struct {
count int
f func() error
errHandler func(error)
predicator RetryPredicator
ansCount int
ansErr error
}
knownErr := errors.New("Duplicate entry '1-389837488-1' for key 'UNI_Trade'")
unknownErr := errors.New("Some Error")
tests := []test{
{
predicator: func(err error) bool {
return !strings.Contains(err.Error(), "Duplicate entry")
},
f: func() error { return knownErr },
ansCount: 1,
ansErr: knownErr,
},
{
predicator: func(err error) bool {
return !strings.Contains(err.Error(), "Duplicate entry")
},
f: func() error { return unknownErr },
ansCount: 3,
ansErr: unknownErr,
},
}
attempts := 3
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
for _, tc := range tests {
err := Retry(ctx, attempts, 100*time.Millisecond, func() error {
tc.count++
return tc.f()
}, tc.errHandler, tc.predicator)
assert.Equal(t, tc.ansCount, tc.count)
assert.EqualError(t, errors.Cause(err), tc.ansErr.Error(), "should be equal")
}
}
func TestRetryCtxCancel(t *testing.T) {
result := int(0)
target := int(3)
ctx, cancel := context.WithCancel(context.Background())
cancel()
err := Retry(ctx, 5, 1*time.Second, func() error { return addAndCheck(&result, target) }, func(error) {})
assert.Error(t, err)
fmt.Println("Error:", err.Error())
assert.Equal(t, int(0), result)
}