Move RateLimiter, ProgressiveRateLimiter, and cleanup functions to ratelimit_shared.go. Add dev and prod rate limiter tests. Remove inline cleanup goroutines. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
194 lines
4.6 KiB
Go
194 lines
4.6 KiB
Go
//go:build test && dev
|
|
// +build test,dev
|
|
|
|
package mw
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"crussell/clock"
|
|
)
|
|
|
|
// ============================================================
|
|
// Dev Stub Behavioral Tests
|
|
// ============================================================
|
|
|
|
// TestDevRateLimiter_AllowAlwaysTrue verifies the dev stub never throttles.
|
|
func TestDevRateLimiter_AllowAlwaysTrue(t *testing.T) {
|
|
rl := NewRateLimiter(10, time.Minute)
|
|
|
|
if !rl.Allow("any-key") {
|
|
t.Error("expected Allow to return true for dev stub")
|
|
}
|
|
if !rl.Allow("another-key") {
|
|
t.Error("expected Allow to return true for dev stub (different key)")
|
|
}
|
|
// Even an empty key should pass
|
|
if !rl.Allow("") {
|
|
t.Error("expected Allow to return true for empty key in dev stub")
|
|
}
|
|
}
|
|
|
|
// TestDevProgressiveRateLimiter_CheckAlwaysZero verifies the dev stub never delays.
|
|
func TestDevProgressiveRateLimiter_CheckAlwaysZero(t *testing.T) {
|
|
prl := NewProgressiveRateLimiter()
|
|
|
|
for i := 0; i < 200; i++ {
|
|
delay := prl.Check("10.0.0.1")
|
|
if delay != 0 {
|
|
t.Errorf("expected 0 delay for dev stub (request %d), got %d", i+1, delay)
|
|
}
|
|
}
|
|
|
|
// Different IP also returns 0
|
|
if delay := prl.Check("10.0.0.2"); delay != 0 {
|
|
t.Errorf("expected 0 delay for different IP in dev stub, got %d", delay)
|
|
}
|
|
}
|
|
|
|
func TestProgressiveRateLimiter_CleanupRemovesStaleEntries(t *testing.T) {
|
|
prl := NewProgressiveRateLimiter()
|
|
|
|
prl.mu.Lock()
|
|
prl.requests["stale-ip"] = &ipProgressiveState{
|
|
timestamps: []time.Time{clock.Now().Add(-120 * time.Second)},
|
|
}
|
|
prl.mu.Unlock()
|
|
|
|
prl.Cleanup()
|
|
|
|
prl.mu.RLock()
|
|
_, exists := prl.requests["stale-ip"]
|
|
prl.mu.RUnlock()
|
|
if exists {
|
|
t.Error("expected stale IP to be cleaned up")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Cleanup(t *testing.T) {
|
|
rl := NewRateLimiter(10, time.Minute)
|
|
|
|
rl.mu.Lock()
|
|
rl.requests["stale-key"] = []time.Time{clock.Now().Add(-5 * time.Minute)}
|
|
rl.requests["fresh-key"] = []time.Time{clock.Now()}
|
|
rl.mu.Unlock()
|
|
|
|
rl.Cleanup()
|
|
|
|
rl.mu.RLock()
|
|
_, staleExists := rl.requests["stale-key"]
|
|
_, freshExists := rl.requests["fresh-key"]
|
|
rl.mu.RUnlock()
|
|
|
|
if staleExists {
|
|
t.Error("expected stale-key to be removed")
|
|
}
|
|
if !freshExists {
|
|
t.Error("expected fresh-key to be preserved")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_Cleanup_EmptyMap(t *testing.T) {
|
|
rl := NewRateLimiter(10, time.Minute)
|
|
|
|
rl.mu.Lock()
|
|
rl.requests = make(map[string][]time.Time)
|
|
rl.mu.Unlock()
|
|
|
|
rl.Cleanup()
|
|
|
|
rl.mu.RLock()
|
|
count := len(rl.requests)
|
|
rl.mu.RUnlock()
|
|
if count != 0 {
|
|
t.Errorf("expected empty map, got %d entries", count)
|
|
}
|
|
}
|
|
|
|
func TestCleanupAllRateLimiters(t *testing.T) {
|
|
registeredLimitersMu.Lock()
|
|
saved := registeredLimiters
|
|
registeredLimitersMu.Unlock()
|
|
defer func() {
|
|
registeredLimitersMu.Lock()
|
|
registeredLimiters = saved
|
|
registeredLimitersMu.Unlock()
|
|
}()
|
|
|
|
rl1 := NewRateLimiter(10, time.Minute)
|
|
rl2 := NewRateLimiter(20, time.Minute)
|
|
|
|
rl1.mu.Lock()
|
|
rl1.requests["rl1-stale"] = []time.Time{clock.Now().Add(-5 * time.Minute)}
|
|
rl1.mu.Unlock()
|
|
|
|
rl2.mu.Lock()
|
|
rl2.requests["rl2-stale"] = []time.Time{clock.Now().Add(-5 * time.Minute)}
|
|
rl2.mu.Unlock()
|
|
|
|
_, err := CleanupAllRateLimiters(context.Background())
|
|
if err != nil {
|
|
t.Errorf("expected nil error, got %v", err)
|
|
}
|
|
|
|
rl1.mu.RLock()
|
|
_, rl1Stale := rl1.requests["rl1-stale"]
|
|
rl1.mu.RUnlock()
|
|
|
|
rl2.mu.RLock()
|
|
_, rl2Stale := rl2.requests["rl2-stale"]
|
|
rl2.mu.RUnlock()
|
|
|
|
if rl1Stale {
|
|
t.Error("expected rl1-stale to be removed")
|
|
}
|
|
if rl2Stale {
|
|
t.Error("expected rl2-stale to be removed")
|
|
}
|
|
}
|
|
|
|
func TestCleanupAllRateLimiters_Empty(t *testing.T) {
|
|
registeredLimitersMu.Lock()
|
|
saved := registeredLimiters
|
|
registeredLimiters = nil
|
|
registeredLimitersMu.Unlock()
|
|
defer func() {
|
|
registeredLimitersMu.Lock()
|
|
registeredLimiters = saved
|
|
registeredLimitersMu.Unlock()
|
|
}()
|
|
|
|
_, err := CleanupAllRateLimiters(context.Background())
|
|
if err != nil {
|
|
t.Errorf("expected nil error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCleanupProgressiveRateLimiter(t *testing.T) {
|
|
globalProgressiveLimiter.mu.Lock()
|
|
saved := globalProgressiveLimiter.requests
|
|
globalProgressiveLimiter.requests = map[string]*ipProgressiveState{
|
|
"global-stale": {timestamps: []time.Time{clock.Now().Add(-120 * time.Second)}},
|
|
}
|
|
globalProgressiveLimiter.mu.Unlock()
|
|
defer func() {
|
|
globalProgressiveLimiter.mu.Lock()
|
|
globalProgressiveLimiter.requests = saved
|
|
globalProgressiveLimiter.mu.Unlock()
|
|
}()
|
|
|
|
_, err := CleanupProgressiveRateLimiter(context.Background())
|
|
if err != nil {
|
|
t.Errorf("expected nil error, got %v", err)
|
|
}
|
|
|
|
globalProgressiveLimiter.mu.RLock()
|
|
_, exists := globalProgressiveLimiter.requests["global-stale"]
|
|
globalProgressiveLimiter.mu.RUnlock()
|
|
if exists {
|
|
t.Error("expected global-stale to be removed")
|
|
}
|
|
}
|