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
|
package ratelimiter
import (
"testing"
"time"
)
func TestPour(t *testing.T) {
bucket := NewLeakyBucket(60, time.Second)
bucket.Lastupdate = time.Unix(0, 0)
bucket.Now = func() time.Time { return time.Unix(1, 0) }
if bucket.Pour(61) {
t.Error("Expected false")
}
if !bucket.Pour(10) {
t.Error("Expected true")
}
if !bucket.Pour(49) {
t.Error("Expected true")
}
if bucket.Pour(2) {
t.Error("Expected false")
}
bucket.Now = func() time.Time { return time.Unix(61, 0) }
if !bucket.Pour(60) {
t.Error("Expected true")
}
if bucket.Pour(1) {
t.Error("Expected false")
}
bucket.Now = func() time.Time { return time.Unix(70, 0) }
if !bucket.Pour(1) {
t.Error("Expected true")
}
}
func TestTimeSinceLastUpdate(t *testing.T) {
bucket := NewLeakyBucket(60, time.Second)
bucket.Now = func() time.Time { return time.Unix(1, 0) }
bucket.Pour(1)
bucket.Now = func() time.Time { return time.Unix(2, 0) }
sinceLast := bucket.TimeSinceLastUpdate()
if sinceLast != time.Second*1 {
t.Errorf("Expected time since last update to be less than 1 second, got %d", sinceLast)
}
}
func TestTimeToDrain(t *testing.T) {
bucket := NewLeakyBucket(60, time.Second)
bucket.Now = func() time.Time { return time.Unix(1, 0) }
bucket.Pour(10)
if bucket.TimeToDrain() != time.Second*10 {
t.Error("Time to drain should be 10 seconds")
}
bucket.Now = func() time.Time { return time.Unix(2, 0) }
if bucket.TimeToDrain() != time.Second*9 {
t.Error("Time to drain should be 9 seconds")
}
}
|
