Add unit tests for RateLimiter (internal/api/ratelimit.go)
16 test cases covering: - Basic rate limiting with Allow() (limit enforcement, different IPs) - Sliding window behavior (time-based expiry, partial expiry) - Edge cases (zero limit, large limit, empty IP) - Cleanup routine (full cleanup, partial cleanup) - Concurrent access (thread safety with 100 parallel attempts) - Middleware integration (allowed/denied paths, X-Forwarded-For) First test file for internal/api package.
This commit is contained in:
rcourtman
parent
c2233fb3f8
commit
722ea83b5d
1 changed files with 433 additions and 0 deletions
433
internal/api/ratelimit_test.go
Normal file
433
internal/api/ratelimit_test.go
Normal file
|
|
@ -0,0 +1,433 @@
|
|||
package api
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestNewRateLimiter(t *testing.T) {
|
||||
rl := NewRateLimiter(5, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
if rl == nil {
|
||||
t.Fatal("NewRateLimiter returned nil")
|
||||
}
|
||||
if rl.limit != 5 {
|
||||
t.Errorf("limit = %d, want 5", rl.limit)
|
||||
}
|
||||
if rl.window != time.Minute {
|
||||
t.Errorf("window = %v, want %v", rl.window, time.Minute)
|
||||
}
|
||||
if rl.attempts == nil {
|
||||
t.Error("attempts map is nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Allow_Basic(t *testing.T) {
|
||||
rl := NewRateLimiter(3, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
ip := "192.168.1.1"
|
||||
|
||||
// First 3 attempts should be allowed
|
||||
for i := 1; i <= 3; i++ {
|
||||
if !rl.Allow(ip) {
|
||||
t.Errorf("attempt %d should be allowed", i)
|
||||
}
|
||||
}
|
||||
|
||||
// 4th attempt should be denied
|
||||
if rl.Allow(ip) {
|
||||
t.Error("4th attempt should be denied")
|
||||
}
|
||||
|
||||
// 5th attempt should also be denied
|
||||
if rl.Allow(ip) {
|
||||
t.Error("5th attempt should be denied")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Allow_DifferentIPs(t *testing.T) {
|
||||
rl := NewRateLimiter(2, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
ip1 := "192.168.1.1"
|
||||
ip2 := "192.168.1.2"
|
||||
|
||||
// Both IPs should have independent limits
|
||||
if !rl.Allow(ip1) {
|
||||
t.Error("ip1 attempt 1 should be allowed")
|
||||
}
|
||||
if !rl.Allow(ip2) {
|
||||
t.Error("ip2 attempt 1 should be allowed")
|
||||
}
|
||||
if !rl.Allow(ip1) {
|
||||
t.Error("ip1 attempt 2 should be allowed")
|
||||
}
|
||||
if !rl.Allow(ip2) {
|
||||
t.Error("ip2 attempt 2 should be allowed")
|
||||
}
|
||||
|
||||
// Both should now be at their limit
|
||||
if rl.Allow(ip1) {
|
||||
t.Error("ip1 attempt 3 should be denied")
|
||||
}
|
||||
if rl.Allow(ip2) {
|
||||
t.Error("ip2 attempt 3 should be denied")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Allow_WindowExpiry(t *testing.T) {
|
||||
// Use a very short window for testing
|
||||
rl := NewRateLimiter(2, 50*time.Millisecond)
|
||||
defer rl.Stop()
|
||||
|
||||
ip := "192.168.1.1"
|
||||
|
||||
// Use up the limit
|
||||
if !rl.Allow(ip) {
|
||||
t.Error("attempt 1 should be allowed")
|
||||
}
|
||||
if !rl.Allow(ip) {
|
||||
t.Error("attempt 2 should be allowed")
|
||||
}
|
||||
if rl.Allow(ip) {
|
||||
t.Error("attempt 3 should be denied")
|
||||
}
|
||||
|
||||
// Wait for window to expire
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
|
||||
// Should be allowed again
|
||||
if !rl.Allow(ip) {
|
||||
t.Error("attempt after window expiry should be allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Allow_SlidingWindow(t *testing.T) {
|
||||
// Test that the window is truly sliding (not fixed intervals)
|
||||
rl := NewRateLimiter(2, 100*time.Millisecond)
|
||||
defer rl.Stop()
|
||||
|
||||
ip := "192.168.1.1"
|
||||
|
||||
// First attempt
|
||||
if !rl.Allow(ip) {
|
||||
t.Error("attempt 1 should be allowed")
|
||||
}
|
||||
|
||||
// Wait half the window
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
|
||||
// Second attempt
|
||||
if !rl.Allow(ip) {
|
||||
t.Error("attempt 2 should be allowed")
|
||||
}
|
||||
|
||||
// Third attempt should be denied (both still in window)
|
||||
if rl.Allow(ip) {
|
||||
t.Error("attempt 3 should be denied")
|
||||
}
|
||||
|
||||
// Wait for first attempt to expire (another 50ms should be enough)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Now should be allowed (first attempt expired, second still valid)
|
||||
if !rl.Allow(ip) {
|
||||
t.Error("attempt 4 should be allowed after first expires")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Allow_ZeroLimit(t *testing.T) {
|
||||
rl := NewRateLimiter(0, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
// With limit 0, nothing should be allowed
|
||||
if rl.Allow("192.168.1.1") {
|
||||
t.Error("should deny with limit 0")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Allow_LargeLimit(t *testing.T) {
|
||||
rl := NewRateLimiter(1000, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
ip := "192.168.1.1"
|
||||
|
||||
// All 1000 attempts should be allowed
|
||||
for i := 1; i <= 1000; i++ {
|
||||
if !rl.Allow(ip) {
|
||||
t.Errorf("attempt %d should be allowed", i)
|
||||
}
|
||||
}
|
||||
|
||||
// 1001st should be denied
|
||||
if rl.Allow(ip) {
|
||||
t.Error("attempt 1001 should be denied")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Allow_EmptyIP(t *testing.T) {
|
||||
rl := NewRateLimiter(2, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
// Empty string IP should still work as a valid key
|
||||
if !rl.Allow("") {
|
||||
t.Error("empty IP attempt 1 should be allowed")
|
||||
}
|
||||
if !rl.Allow("") {
|
||||
t.Error("empty IP attempt 2 should be allowed")
|
||||
}
|
||||
if rl.Allow("") {
|
||||
t.Error("empty IP attempt 3 should be denied")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Cleanup(t *testing.T) {
|
||||
rl := NewRateLimiter(5, 50*time.Millisecond)
|
||||
defer rl.Stop()
|
||||
|
||||
// Add some attempts
|
||||
rl.Allow("192.168.1.1")
|
||||
rl.Allow("192.168.1.2")
|
||||
rl.Allow("192.168.1.3")
|
||||
|
||||
// Verify attempts are tracked
|
||||
rl.mu.RLock()
|
||||
beforeCount := len(rl.attempts)
|
||||
rl.mu.RUnlock()
|
||||
|
||||
if beforeCount != 3 {
|
||||
t.Errorf("expected 3 IPs tracked, got %d", beforeCount)
|
||||
}
|
||||
|
||||
// Wait for window to expire
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
|
||||
// Manually trigger cleanup
|
||||
rl.cleanup()
|
||||
|
||||
// All should be cleaned up
|
||||
rl.mu.RLock()
|
||||
afterCount := len(rl.attempts)
|
||||
rl.mu.RUnlock()
|
||||
|
||||
if afterCount != 0 {
|
||||
t.Errorf("expected 0 IPs after cleanup, got %d", afterCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Cleanup_PartialExpiry(t *testing.T) {
|
||||
rl := NewRateLimiter(5, 50*time.Millisecond)
|
||||
defer rl.Stop()
|
||||
|
||||
// Add attempt for IP1
|
||||
rl.Allow("192.168.1.1")
|
||||
|
||||
// Wait for it to expire
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
|
||||
// Add attempt for IP2 (this one is fresh)
|
||||
rl.Allow("192.168.1.2")
|
||||
|
||||
// Cleanup should remove IP1 but keep IP2
|
||||
rl.cleanup()
|
||||
|
||||
rl.mu.RLock()
|
||||
_, hasIP1 := rl.attempts["192.168.1.1"]
|
||||
_, hasIP2 := rl.attempts["192.168.1.2"]
|
||||
rl.mu.RUnlock()
|
||||
|
||||
if hasIP1 {
|
||||
t.Error("IP1 should have been cleaned up")
|
||||
}
|
||||
if !hasIP2 {
|
||||
t.Error("IP2 should still be present")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Concurrent(t *testing.T) {
|
||||
rl := NewRateLimiter(100, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 10
|
||||
attemptsPerGoroutine := 20
|
||||
|
||||
// Track results per goroutine
|
||||
results := make([][]bool, numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
results[i] = make([]bool, attemptsPerGoroutine)
|
||||
}
|
||||
|
||||
// Concurrent access from same IP
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
for j := 0; j < attemptsPerGoroutine; j++ {
|
||||
results[idx][j] = rl.Allow("shared-ip")
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Count allowed attempts
|
||||
allowed := 0
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
for j := 0; j < attemptsPerGoroutine; j++ {
|
||||
if results[i][j] {
|
||||
allowed++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Should have exactly 100 allowed (the limit)
|
||||
if allowed != 100 {
|
||||
t.Errorf("expected exactly 100 allowed attempts, got %d", allowed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Stop(t *testing.T) {
|
||||
rl := NewRateLimiter(5, time.Minute)
|
||||
|
||||
// Should be able to stop without issues
|
||||
rl.Stop()
|
||||
|
||||
// After stop, Allow should still work (stop only affects cleanup goroutine)
|
||||
if !rl.Allow("192.168.1.1") {
|
||||
t.Error("Allow should still work after Stop")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Middleware_Allowed(t *testing.T) {
|
||||
rl := NewRateLimiter(5, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
handlerCalled := false
|
||||
handler := func(w http.ResponseWriter, r *http.Request) {
|
||||
handlerCalled = true
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
|
||||
wrapped := rl.Middleware(handler)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.1:12345"
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
wrapped(w, req)
|
||||
|
||||
if !handlerCalled {
|
||||
t.Error("handler should have been called")
|
||||
}
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("status = %d, want %d", w.Code, http.StatusOK)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Middleware_Denied(t *testing.T) {
|
||||
rl := NewRateLimiter(1, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
handlerCalls := 0
|
||||
handler := func(w http.ResponseWriter, r *http.Request) {
|
||||
handlerCalls++
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}
|
||||
|
||||
wrapped := rl.Middleware(handler)
|
||||
|
||||
// First request should succeed
|
||||
req1 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req1.RemoteAddr = "192.168.1.1:12345"
|
||||
w1 := httptest.NewRecorder()
|
||||
wrapped(w1, req1)
|
||||
|
||||
if w1.Code != http.StatusOK {
|
||||
t.Errorf("first request status = %d, want %d", w1.Code, http.StatusOK)
|
||||
}
|
||||
|
||||
// Second request should be rate limited
|
||||
req2 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req2.RemoteAddr = "192.168.1.1:12345"
|
||||
w2 := httptest.NewRecorder()
|
||||
wrapped(w2, req2)
|
||||
|
||||
if w2.Code != http.StatusTooManyRequests {
|
||||
t.Errorf("second request status = %d, want %d", w2.Code, http.StatusTooManyRequests)
|
||||
}
|
||||
|
||||
if handlerCalls != 1 {
|
||||
t.Errorf("handler called %d times, want 1", handlerCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Middleware_XForwardedFor(t *testing.T) {
|
||||
rl := NewRateLimiter(1, time.Minute)
|
||||
defer rl.Stop()
|
||||
|
||||
handlerCalls := 0
|
||||
handler := func(w http.ResponseWriter, r *http.Request) {
|
||||
handlerCalls++
|
||||
}
|
||||
|
||||
wrapped := rl.Middleware(handler)
|
||||
|
||||
// Request with X-Forwarded-For header
|
||||
req1 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req1.RemoteAddr = "127.0.0.1:12345"
|
||||
req1.Header.Set("X-Forwarded-For", "203.0.113.1")
|
||||
w1 := httptest.NewRecorder()
|
||||
wrapped(w1, req1)
|
||||
|
||||
// Second request with same X-Forwarded-For should be rate limited
|
||||
req2 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req2.RemoteAddr = "127.0.0.1:12345"
|
||||
req2.Header.Set("X-Forwarded-For", "203.0.113.1")
|
||||
w2 := httptest.NewRecorder()
|
||||
wrapped(w2, req2)
|
||||
|
||||
if w2.Code != http.StatusTooManyRequests {
|
||||
t.Errorf("second request with same X-Forwarded-For should be rate limited, got %d", w2.Code)
|
||||
}
|
||||
|
||||
// Request with different X-Forwarded-For should be allowed
|
||||
req3 := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req3.RemoteAddr = "127.0.0.1:12345"
|
||||
req3.Header.Set("X-Forwarded-For", "203.0.113.2")
|
||||
w3 := httptest.NewRecorder()
|
||||
wrapped(w3, req3)
|
||||
|
||||
if w3.Code == http.StatusTooManyRequests {
|
||||
t.Error("request with different X-Forwarded-For should be allowed")
|
||||
}
|
||||
|
||||
if handlerCalls != 2 {
|
||||
t.Errorf("handler called %d times, want 2", handlerCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimiter_Middleware_ResponseBody(t *testing.T) {
|
||||
rl := NewRateLimiter(0, time.Minute) // Immediately rate limit
|
||||
defer rl.Stop()
|
||||
|
||||
handler := func(w http.ResponseWriter, r *http.Request) {}
|
||||
wrapped := rl.Middleware(handler)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.RemoteAddr = "192.168.1.1:12345"
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
wrapped(w, req)
|
||||
|
||||
body := w.Body.String()
|
||||
expected := "Rate limit exceeded. Please try again later.\n"
|
||||
if body != expected {
|
||||
t.Errorf("body = %q, want %q", body, expected)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue