/
rate_test.go
128 lines (108 loc) · 3.74 KB
/
rate_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
package throttled_test
import (
"testing"
"time"
"github.com/throttled/throttled"
"github.com/throttled/throttled/store/memstore"
)
const deniedStatus = 429
type testStore struct {
store throttled.GCRAStore
clock time.Time
failUpdates bool
}
func (ts *testStore) GetWithTime(key string) (int64, time.Time, error) {
v, _, e := ts.store.GetWithTime(key)
return v, ts.clock, e
}
func (ts *testStore) SetIfNotExistsWithTTL(key string, value int64, ttl time.Duration) (bool, error) {
if ts.failUpdates {
return false, nil
}
return ts.store.SetIfNotExistsWithTTL(key, value, ttl)
}
func (ts *testStore) CompareAndSwapWithTTL(key string, old, new int64, ttl time.Duration) (bool, error) {
if ts.failUpdates {
return false, nil
}
return ts.store.CompareAndSwapWithTTL(key, old, new, ttl)
}
func TestRateLimit(t *testing.T) {
limit := 5
rq := throttled.RateQuota{MaxRate: throttled.PerSec(1), MaxBurst: limit - 1}
start := time.Unix(0, 0)
cases := []struct {
now time.Time
volume, remaining int
reset, retry time.Duration
limited bool
}{
// You can never make a request larger than the maximum
0: {start, 6, 5, 0, -1, true},
// Rate limit normal requests appropriately
1: {start, 1, 4, time.Second, -1, false},
2: {start, 1, 3, 2 * time.Second, -1, false},
3: {start, 1, 2, 3 * time.Second, -1, false},
4: {start, 1, 1, 4 * time.Second, -1, false},
5: {start, 1, 0, 5 * time.Second, -1, false},
6: {start, 1, 0, 5 * time.Second, time.Second, true},
7: {start.Add(3000 * time.Millisecond), 1, 2, 3000 * time.Millisecond, -1, false},
8: {start.Add(3100 * time.Millisecond), 1, 1, 3900 * time.Millisecond, -1, false},
9: {start.Add(4000 * time.Millisecond), 1, 1, 4000 * time.Millisecond, -1, false},
10: {start.Add(8000 * time.Millisecond), 1, 4, 1000 * time.Millisecond, -1, false},
11: {start.Add(9500 * time.Millisecond), 1, 4, 1000 * time.Millisecond, -1, false},
// Zero-volume request just peeks at the state
12: {start.Add(9500 * time.Millisecond), 0, 4, time.Second, -1, false},
// High-volume request uses up more of the limit
13: {start.Add(9500 * time.Millisecond), 2, 2, 3 * time.Second, -1, false},
// Large requests cannot exceed limits
14: {start.Add(9500 * time.Millisecond), 5, 2, 3 * time.Second, 3 * time.Second, true},
}
mst, err := memstore.New(0)
if err != nil {
t.Fatal(err)
}
st := testStore{store: mst}
rl, err := throttled.NewGCRARateLimiter(&st, rq)
if err != nil {
t.Fatal(err)
}
// Start the server
for i, c := range cases {
st.clock = c.now
limited, context, err := rl.RateLimit("foo", c.volume)
if err != nil {
t.Fatalf("%d: %#v", i, err)
}
if limited != c.limited {
t.Errorf("%d: expected Limited to be %t but got %t", i, c.limited, limited)
}
if have, want := context.Limit, limit; have != want {
t.Errorf("%d: expected Limit to be %d but got %d", i, want, have)
}
if have, want := context.Remaining, c.remaining; have != want {
t.Errorf("%d: expected Remaining to be %d but got %d", i, want, have)
}
if have, want := context.ResetAfter, c.reset; have != want {
t.Errorf("%d: expected ResetAfter to be %s but got %s", i, want, have)
}
if have, want := context.RetryAfter, c.retry; have != want {
t.Errorf("%d: expected RetryAfter to be %d but got %d", i, want, have)
}
}
}
func TestRateLimitUpdateFailures(t *testing.T) {
rq := throttled.RateQuota{MaxRate: throttled.PerSec(1), MaxBurst: 1}
mst, err := memstore.New(0)
if err != nil {
t.Fatal(err)
}
st := testStore{store: mst, failUpdates: true}
rl, err := throttled.NewGCRARateLimiter(&st, rq)
if err != nil {
t.Fatal(err)
}
if _, _, err := rl.RateLimit("foo", 1); err == nil {
t.Error("Expected limiting to fail when store updates fail")
}
}