9 "github.com/go-kit/kit/endpoint"
10 "github.com/go-kit/kit/sd"
11 "github.com/go-kit/kit/sd/lb"
14 func TestRetryMaxTotalFail(t *testing.T) {
16 endpoints = sd.FixedEndpointer{} // no endpoints
17 rr = lb.NewRoundRobin(endpoints)
18 retry = lb.Retry(999, time.Second, rr) // lots of retries
19 ctx = context.Background()
21 if _, err := retry(ctx, struct{}{}); err == nil {
22 t.Errorf("expected error, got none") // should fail
26 func TestRetryMaxPartialFail(t *testing.T) {
28 endpoints = []endpoint.Endpoint{
29 func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error one") },
30 func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error two") },
31 func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
33 endpointer = sd.FixedEndpointer{
38 retries = len(endpoints) - 1 // not quite enough retries
39 rr = lb.NewRoundRobin(endpointer)
40 ctx = context.Background()
42 if _, err := lb.Retry(retries, time.Second, rr)(ctx, struct{}{}); err == nil {
43 t.Errorf("expected error two, got none")
47 func TestRetryMaxSuccess(t *testing.T) {
49 endpoints = []endpoint.Endpoint{
50 func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error one") },
51 func(context.Context, interface{}) (interface{}, error) { return nil, errors.New("error two") },
52 func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
54 endpointer = sd.FixedEndpointer{
59 retries = len(endpoints) // exactly enough retries
60 rr = lb.NewRoundRobin(endpointer)
61 ctx = context.Background()
63 if _, err := lb.Retry(retries, time.Second, rr)(ctx, struct{}{}); err != nil {
68 func TestRetryTimeout(t *testing.T) {
70 step = make(chan struct{})
71 e = func(context.Context, interface{}) (interface{}, error) { <-step; return struct{}{}, nil }
72 timeout = time.Millisecond
73 retry = lb.Retry(999, timeout, lb.NewRoundRobin(sd.FixedEndpointer{0: e}))
74 errs = make(chan error, 1)
75 invoke = func() { _, err := retry(context.Background(), struct{}{}); errs <- err }
78 go func() { step <- struct{}{} }() // queue up a flush of the endpoint
79 invoke() // invoke the endpoint and trigger the flush
80 if err := <-errs; err != nil { // that should succeed
84 go func() { time.Sleep(10 * timeout); step <- struct{}{} }() // a delayed flush
85 invoke() // invoke the endpoint
86 if err := <-errs; err != context.DeadlineExceeded { // that should not succeed
87 t.Errorf("wanted %v, got none", context.DeadlineExceeded)
91 func TestAbortEarlyCustomMessage(t *testing.T) {
93 myErr = errors.New("aborting early")
94 cb = func(int, error) (bool, error) { return false, myErr }
95 endpoints = sd.FixedEndpointer{} // no endpoints
96 rr = lb.NewRoundRobin(endpoints)
97 retry = lb.RetryWithCallback(time.Second, rr, cb) // lots of retries
98 ctx = context.Background()
100 _, err := retry(ctx, struct{}{})
101 if want, have := myErr, err.(lb.RetryError).Final; want != have {
102 t.Errorf("want %v, have %v", want, have)
106 func TestErrorPassedUnchangedToCallback(t *testing.T) {
108 myErr = errors.New("my custom error")
109 cb = func(_ int, err error) (bool, error) {
110 if want, have := myErr, err; want != have {
111 t.Errorf("want %v, have %v", want, have)
115 endpoint = func(ctx context.Context, request interface{}) (interface{}, error) {
118 endpoints = sd.FixedEndpointer{endpoint} // no endpoints
119 rr = lb.NewRoundRobin(endpoints)
120 retry = lb.RetryWithCallback(time.Second, rr, cb) // lots of retries
121 ctx = context.Background()
123 _, err := retry(ctx, struct{}{})
124 if want, have := myErr, err.(lb.RetryError).Final; want != have {
125 t.Errorf("want %v, have %v", want, have)
129 func TestHandleNilCallback(t *testing.T) {
131 endpointer = sd.FixedEndpointer{
132 func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil /* OK */ },
134 rr = lb.NewRoundRobin(endpointer)
135 ctx = context.Background()
137 retry := lb.RetryWithCallback(time.Second, rr, nil)
138 if _, err := retry(ctx, struct{}{}); err != nil {