1

internal/ratelimit/ratelimit_test.go

@@ -0,0 +1,200 @@
+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)
+	}
+}