package ratelimit import ( "context" "testing" "time" ) func TestTokenBucket_Allow(t *testing.T) { bucket := NewTokenBucket(5, 1) // 5 capacity, 1 token per second // Should allow 5 requests initially for i := 0; i < 5; i++ { if !bucket.Allow() { t.Errorf("Request %d should be allowed", i+1) } } // 6th request should be denied if bucket.Allow() { t.Error("6th request should be denied") } } func TestTokenBucket_Refill(t *testing.T) { bucket := NewTokenBucket(2, 2) // 2 capacity, 2 tokens per second // Consume all tokens bucket.Allow() bucket.Allow() // Should be empty now if bucket.Allow() { t.Error("Bucket should be empty") } // Wait for refill time.Sleep(1100 * time.Millisecond) // Wait a bit more than 1 second // Should have tokens again if !bucket.Allow() { t.Error("Should have tokens after refill") } } func TestTokenBucket_Wait(t *testing.T) { bucket := NewTokenBucket(1, 1) // 1 capacity, 1 token per second // Consume the token if !bucket.Allow() { t.Error("First request should be allowed") } // Test wait with timeout ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() start := time.Now() err := bucket.Wait(ctx) duration := time.Since(start) if err != context.DeadlineExceeded { t.Errorf("Expected context.DeadlineExceeded, got %v", err) } if duration < 90*time.Millisecond || duration > 150*time.Millisecond { t.Errorf("Wait duration should be around 100ms, got %v", duration) } } func TestFixedWindowLimiter_Allow(t *testing.T) { limiter := NewFixedWindowLimiter(3, 1*time.Second) // 3 requests per second // Should allow 3 requests for i := 0; i < 3; i++ { if !limiter.Allow() { t.Errorf("Request %d should be allowed", i+1) } } // 4th request should be denied if limiter.Allow() { t.Error("4th request should be denied") } } func TestFixedWindowLimiter_WindowReset(t *testing.T) { limiter := NewFixedWindowLimiter(2, 500*time.Millisecond) // 2 requests per 500ms // Consume all requests limiter.Allow() limiter.Allow() // Should be at limit if limiter.Allow() { t.Error("Should be at limit") } // Wait for window reset time.Sleep(600 * time.Millisecond) // Should allow requests again if !limiter.Allow() { t.Error("Should allow requests after window reset") } } func TestAdaptiveLimiter_Basic(t *testing.T) { limiter := NewAdaptiveLimiter(10, 100*time.Millisecond, 0.05) // 10 RPS, 100ms target, 5% max error // Should allow initial requests for i := 0; i < 5; i++ { if !limiter.Allow() { t.Errorf("Request %d should be allowed", i+1) } } } func TestAdaptiveLimiter_RecordResponse(t *testing.T) { limiter := NewAdaptiveLimiter(10, 100*time.Millisecond, 0.05) // Record some fast responses for i := 0; i < 15; i++ { limiter.RecordResponse(50*time.Millisecond, true) } // Check that rate might have increased (or at least not decreased significantly) initialRate := limiter.GetCurrentRate() if initialRate < 8 { // Should be at least close to original rate t.Errorf("Rate should not decrease significantly with good responses, got %d", initialRate) } } func TestAdaptiveLimiter_SlowResponses(t *testing.T) { limiter := NewAdaptiveLimiter(10, 100*time.Millisecond, 0.05) // Record some slow responses for i := 0; i < 15; i++ { limiter.RecordResponse(500*time.Millisecond, true) // 5x target latency } // Rate should decrease finalRate := limiter.GetCurrentRate() if finalRate >= 10 { t.Errorf("Rate should decrease with slow responses, got %d", finalRate) } } func TestAdaptiveLimiter_HighErrorRate(t *testing.T) { limiter := NewAdaptiveLimiter(10, 100*time.Millisecond, 0.05) // Record responses with high error rate for i := 0; i < 15; i++ { success := i < 5 // Only first 5 are successful (33% success rate) limiter.RecordResponse(50*time.Millisecond, success) } // Rate should decrease due to high error rate finalRate := limiter.GetCurrentRate() if finalRate >= 10 { t.Errorf("Rate should decrease with high error rate, got %d", finalRate) } } // Benchmark tests func BenchmarkTokenBucket_Allow(b *testing.B) { bucket := NewTokenBucket(1000, 1000) b.ResetTimer() for i := 0; i < b.N; i++ { bucket.Allow() } } func BenchmarkFixedWindowLimiter_Allow(b *testing.B) { limiter := NewFixedWindowLimiter(1000, 1*time.Second) b.ResetTimer() for i := 0; i < b.N; i++ { limiter.Allow() } } func BenchmarkAdaptiveLimiter_Allow(b *testing.B) { limiter := NewAdaptiveLimiter(1000, 100*time.Millisecond, 0.05) b.ResetTimer() for i := 0; i < b.N; i++ { limiter.Allow() } } func BenchmarkAdaptiveLimiter_RecordResponse(b *testing.B) { limiter := NewAdaptiveLimiter(1000, 100*time.Millisecond, 0.05) b.ResetTimer() for i := 0; i < b.N; i++ { limiter.RecordResponse(50*time.Millisecond, true) } }