internal/control/ratelimit_test.go

 1package control
 2
 3import (
 4	"testing"
 5	"time"
 6)
 7
 8// The budget is per account: spending one account's does not touch
 9// another's, which is the property that makes the limit meaningful when
10// registration is open (#148).
11func TestWriteLimiterIsPerAccount(t *testing.T) {
12	l := &writeLimiter{buckets: map[int64]*writeBucket{}}
13	for i := 0; i < 3; i++ {
14		if ok, _ := l.allow(1, 3); !ok {
15			t.Fatalf("write %d of the burst refused", i+1)
16		}
17	}
18	if ok, wait := l.allow(1, 3); ok || wait <= 0 {
19		t.Errorf("the fourth write was allowed, or gave no wait: %v", wait)
20	}
21	if ok, _ := l.allow(2, 3); !ok {
22		t.Error("a second account was charged for the first's writes")
23	}
24}
25
26// An idle account refills at the sustained rate, so the limit throttles
27// rather than locks out.
28func TestWriteLimiterRefills(t *testing.T) {
29	l := &writeLimiter{buckets: map[int64]*writeBucket{}}
30	if ok, _ := l.allow(1, 60); !ok {
31		t.Fatal("first write refused")
32	}
33	// 60 a minute is one a second: rewind the clock a second and the
34	// spent token is back.
35	l.buckets[1].tokens = 0
36	l.buckets[1].last = time.Now().Add(-time.Second)
37	if ok, wait := l.allow(1, 60); !ok {
38		t.Errorf("no refill after a second: wait %v", wait)
39	}
40}
41
42// The wait a refusal reports is the time until a token exists, so a
43// caller that honours it is not refused again immediately.
44func TestWriteLimiterReportsAUsableWait(t *testing.T) {
45	l := &writeLimiter{buckets: map[int64]*writeBucket{}}
46	l.allow(1, 60)
47	l.buckets[1].tokens = 0
48	ok, wait := l.allow(1, 60)
49	if ok {
50		t.Fatal("allowed with an empty bucket")
51	}
52	if wait < time.Second || wait > 3*time.Second {
53		t.Errorf("wait %v is not the ~1s a 60/min bucket needs", wait)
54	}
55}