1 | 1 |
|
2 | 2 |
import (
|
3 | 3 |
"context"
|
4 | |
"math"
|
5 | 4 |
"strings"
|
6 | 5 |
"testing"
|
7 | 6 |
"time"
|
8 | 7 |
|
9 | 8 |
jujuratelimit "github.com/juju/ratelimit"
|
|
9 |
"golang.org/x/time/rate"
|
10 | 10 |
|
11 | 11 |
"github.com/go-kit/kit/endpoint"
|
12 | 12 |
"github.com/go-kit/kit/ratelimit"
|
13 | |
"golang.org/x/time/rate"
|
14 | 13 |
)
|
15 | 14 |
|
|
15 |
var nopEndpoint = func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
|
|
16 |
|
16 | 17 |
func TestTokenBucketLimiter(t *testing.T) {
|
17 | |
e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
|
18 | |
for _, n := range []int{1, 2, 100} {
|
19 | |
tb := jujuratelimit.NewBucketWithRate(float64(n), int64(n))
|
20 | |
testLimiter(t, ratelimit.NewTokenBucketLimiter(tb)(e), n)
|
21 | |
}
|
|
18 |
tb := jujuratelimit.NewBucket(time.Minute, 1)
|
|
19 |
testSuccessThenFailure(
|
|
20 |
t,
|
|
21 |
ratelimit.NewTokenBucketLimiter(tb)(nopEndpoint),
|
|
22 |
ratelimit.ErrLimited.Error())
|
22 | 23 |
}
|
23 | 24 |
|
24 | 25 |
func TestTokenBucketThrottler(t *testing.T) {
|
25 | |
d := time.Duration(0)
|
26 | |
s := func(d0 time.Duration) { d = d0 }
|
27 | |
|
28 | |
e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
|
29 | |
e = ratelimit.NewTokenBucketThrottler(jujuratelimit.NewBucketWithRate(1, 1), s)(e)
|
30 | |
|
31 | |
// First request should go through with no delay.
|
32 | |
e(context.Background(), struct{}{})
|
33 | |
if want, have := time.Duration(0), d; want != have {
|
34 | |
t.Errorf("want %s, have %s", want, have)
|
35 | |
}
|
36 | |
|
37 | |
// Next request should request a ~1s sleep.
|
38 | |
e(context.Background(), struct{}{})
|
39 | |
if want, have, tol := time.Second, d, time.Millisecond; math.Abs(float64(want-have)) > float64(tol) {
|
40 | |
t.Errorf("want %s, have %s", want, have)
|
41 | |
}
|
42 | |
}
|
43 | |
|
44 | |
func testLimiter(t *testing.T, e endpoint.Endpoint, rate int) {
|
45 | |
// First <rate> requests should succeed.
|
46 | |
for i := 0; i < rate; i++ {
|
47 | |
if _, err := e(context.Background(), struct{}{}); err != nil {
|
48 | |
t.Fatalf("rate=%d: request %d/%d failed: %v", rate, i+1, rate, err)
|
49 | |
}
|
50 | |
}
|
51 | |
|
52 | |
// Next request should fail.
|
53 | |
if _, err := e(context.Background(), struct{}{}); err != ratelimit.ErrLimited {
|
54 | |
t.Errorf("rate=%d: want %v, have %v", rate, ratelimit.ErrLimited, err)
|
55 | |
}
|
|
26 |
tb := jujuratelimit.NewBucket(time.Minute, 1)
|
|
27 |
testSuccessThenFailure(
|
|
28 |
t,
|
|
29 |
ratelimit.NewTokenBucketThrottler(tb, nil)(nopEndpoint),
|
|
30 |
"context deadline exceeded")
|
56 | 31 |
}
|
57 | 32 |
|
58 | 33 |
func TestXRateErroring(t *testing.T) {
|
59 | 34 |
limit := rate.NewLimiter(rate.Every(time.Minute), 1)
|
60 | |
e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
|
61 | |
testLimiter(t, ratelimit.NewErroringLimiter(limit)(e), 1)
|
|
35 |
testSuccessThenFailure(
|
|
36 |
t,
|
|
37 |
ratelimit.NewErroringLimiter(limit)(nopEndpoint),
|
|
38 |
ratelimit.ErrLimited.Error())
|
62 | 39 |
}
|
63 | 40 |
|
64 | 41 |
func TestXRateDelaying(t *testing.T) {
|
65 | 42 |
limit := rate.NewLimiter(rate.Every(time.Minute), 1)
|
66 | |
e := func(context.Context, interface{}) (interface{}, error) { return struct{}{}, nil }
|
67 | |
e = ratelimit.NewDelayingLimiter(limit)(e)
|
|
43 |
testSuccessThenFailure(
|
|
44 |
t,
|
|
45 |
ratelimit.NewDelayingLimiter(limit)(nopEndpoint),
|
|
46 |
"exceed context deadline")
|
|
47 |
}
|
68 | 48 |
|
69 | |
_, err := e(context.Background(), struct{}{})
|
70 | |
if err != nil {
|
|
49 |
func testSuccessThenFailure(t *testing.T, e endpoint.Endpoint, failContains string) {
|
|
50 |
ctx, cxl := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
|
51 |
defer cxl()
|
|
52 |
|
|
53 |
// First request should succeed.
|
|
54 |
if _, err := e(ctx, struct{}{}); err != nil {
|
71 | 55 |
t.Errorf("unexpected: %v\n", err)
|
72 | 56 |
}
|
73 | 57 |
|
74 | |
dur := 500 * time.Millisecond
|
75 | |
ctx, cxl := context.WithTimeout(context.Background(), dur)
|
76 | |
defer cxl()
|
77 | |
|
78 | |
_, err = e(ctx, struct{}{})
|
79 | |
if !strings.Contains(err.Error(), "exceed context deadline") {
|
80 | |
t.Errorf("expected timeout: %v\n", err)
|
|
58 |
// Next request should fail.
|
|
59 |
if _, err := e(ctx, struct{}{}); !strings.Contains(err.Error(), failContains) {
|
|
60 |
t.Errorf("expected `%s`: %v\n", failContains, err)
|
81 | 61 |
}
|
82 | 62 |
}
|