1 package ratelimit
2 3 import (
4 "testing"
5 "time"
6 )
7 8 // TestBurstThenDeny is the property the relay depends on: a key gets exactly
9 // burst writes, and the state has to survive between calls.
10 func TestBurstThenDeny(t *testing.T) {
11 l := New(0.000001, 3)
12 key := []byte("1.2.3.4")
13 for i := int32(0); i < 3; i++ {
14 if !l.Allow(key) {
15 t.Fatalf("token %d should be allowed", i)
16 }
17 }
18 if l.Allow(key) {
19 t.Fatal("burst is exhausted, the next write must be denied")
20 }
21 if len(l.buckets) != 1 {
22 t.Fatalf("one bucket expected, got %d", len(l.buckets))
23 }
24 }
25 26 func TestDistinctKeys(t *testing.T) {
27 l := New(0.000001, 1)
28 if !l.Allow([]byte("a")) {
29 t.Fatal("first key should be allowed")
30 }
31 if l.Allow([]byte("a")) {
32 t.Fatal("first key is spent")
33 }
34 if !l.Allow([]byte("b")) {
35 t.Fatal("a second key has its own burst")
36 }
37 }
38 39 func TestRefill(t *testing.T) {
40 // 10 tokens/second: one token needs 100ms, so the immediate second call
41 // cannot have refilled.
42 l := New(10.0, 1)
43 key := []byte("k")
44 if !l.Allow(key) {
45 t.Fatal("first allow")
46 }
47 if l.Allow(key) {
48 t.Fatal("no token should be left")
49 }
50 time.Sleep(200 * time.Millisecond)
51 if !l.Allow(key) {
52 t.Fatal("the bucket should have refilled")
53 }
54 }
55 56 // TestRefillIsCapped checks that a long idle period does not accumulate more
57 // than burst tokens.
58 func TestRefillIsCapped(t *testing.T) {
59 l := New(10.0, 2)
60 key := []byte("k")
61 time.Sleep(300 * time.Millisecond)
62 for i := int32(0); i < 2; i++ {
63 if !l.Allow(key) {
64 t.Fatalf("token %d should be allowed", i)
65 }
66 }
67 if l.Allow(key) {
68 t.Fatal("burst is 2, a third token must not be granted")
69 }
70 }
71 72 func TestCleanup(t *testing.T) {
73 l := New(1.0, 1)
74 if !l.Allow([]byte("old")) {
75 t.Fatal("allow")
76 }
77 if len(l.buckets) != 1 {
78 t.Fatal("bucket should be recorded")
79 }
80 l.Cleanup(time.Duration(0))
81 if len(l.buckets) != 0 {
82 t.Fatal("Cleanup(0) should drop every bucket")
83 }
84 // The limiter still works after its state was dropped.
85 if !l.Allow([]byte("new")) {
86 t.Fatal("allow after cleanup")
87 }
88 }
89