stroke/internal/ratelimit/ratelimit_test.go

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)
	}
}