102 lines
2.6 KiB
Go
102 lines
2.6 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestLimiter_AllowsUnderThreshold(t *testing.T) {
|
|
l := NewLimiter(10, time.Minute)
|
|
for i := range 9 {
|
|
if !l.Allowed("1.2.3.4") {
|
|
t.Fatalf("expected allowed at failure %d", i)
|
|
}
|
|
l.RecordFailure("1.2.3.4")
|
|
}
|
|
if !l.Allowed("1.2.3.4") {
|
|
t.Fatalf("expected allowed after 9 failures")
|
|
}
|
|
}
|
|
|
|
func TestLimiter_BlocksAtThreshold(t *testing.T) {
|
|
l := NewLimiter(10, time.Minute)
|
|
for range 10 {
|
|
l.RecordFailure("1.2.3.4")
|
|
}
|
|
if l.Allowed("1.2.3.4") {
|
|
t.Fatalf("expected blocked after 10 failures")
|
|
}
|
|
}
|
|
|
|
func TestLimiter_BlockedClientStaysThrottledAcrossWindowBoundary(t *testing.T) {
|
|
window := 300 * time.Millisecond
|
|
l := NewLimiter(10, window)
|
|
|
|
for range 20 {
|
|
l.RecordFailure("1.2.3.4")
|
|
}
|
|
if l.Allowed("1.2.3.4") {
|
|
t.Fatalf("expected blocked immediately after 20 failures")
|
|
}
|
|
|
|
idx := l.windowIndex(time.Now())
|
|
windowEnd := time.Unix(0, (idx+1)*int64(window))
|
|
// Sleep just past the window boundary so the failures roll into the previous
|
|
// bucket, then check that their weighted contribution still blocks the client.
|
|
time.Sleep(time.Until(windowEnd) + window/10)
|
|
|
|
if l.Allowed("1.2.3.4") {
|
|
t.Fatalf("expected still blocked shortly after the window boundary")
|
|
}
|
|
}
|
|
|
|
func TestLimiter_ResetsAfterIdlePeriod(t *testing.T) {
|
|
window := 50 * time.Millisecond
|
|
l := NewLimiter(10, window)
|
|
|
|
for range 10 {
|
|
l.RecordFailure("1.2.3.4")
|
|
}
|
|
if l.Allowed("1.2.3.4") {
|
|
t.Fatalf("expected blocked immediately after 10 failures")
|
|
}
|
|
|
|
time.Sleep(window * 3)
|
|
|
|
if !l.Allowed("1.2.3.4") {
|
|
t.Fatalf("expected allowed after idling for 2+ windows")
|
|
}
|
|
}
|
|
|
|
func TestLimiter_Cleanup(t *testing.T) {
|
|
window := 20 * time.Millisecond
|
|
l := NewLimiter(10, window)
|
|
l.RecordFailure("1.2.3.4")
|
|
|
|
time.Sleep(window * 4)
|
|
l.Cleanup()
|
|
|
|
l.mu.RLock()
|
|
_, ok := l.state["1.2.3.4"]
|
|
l.mu.RUnlock()
|
|
if ok {
|
|
t.Fatalf("expected stale IP to be cleaned up")
|
|
}
|
|
}
|
|
|
|
func TestOverlapFraction(t *testing.T) {
|
|
window := 100 * time.Millisecond
|
|
previousIdx := int64(1000)
|
|
previousWindowEnd := time.Unix(0, (previousIdx+1)*int64(window))
|
|
|
|
if got := overlapFraction(previousWindowEnd.Add(-10*time.Millisecond), window, previousIdx); got != 1 {
|
|
t.Fatalf("expected overlap 1 before window end, got %v", got)
|
|
}
|
|
if got := overlapFraction(previousWindowEnd.Add(20*time.Millisecond), window, previousIdx); got < 0.79 || got > 0.81 {
|
|
t.Fatalf("expected overlap ~0.8 at 20ms into next window, got %v", got)
|
|
}
|
|
if got := overlapFraction(previousWindowEnd.Add(200*time.Millisecond), window, previousIdx); got != 0 {
|
|
t.Fatalf("expected overlap 0 once fully elapsed, got %v", got)
|
|
}
|
|
}
|