ratelimiter-test.go

package ratelimiter

import (
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestTokenBucket(t *testing.T) {
	// Create limiter: 10 tokens/sec, capacity 20
	tb := NewTokenBucket(10, 20)
	key := "test-user"

	// Test 1: Initial burst capacity
	for i := 0; i < 20; i++ {
		allowed, err := tb.Allow(key)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if !allowed {
			t.Fatalf("expected request %d to be allowed", i+1)
		}
	}

	// Test 2: Should be rate limited now
	allowed, _ := tb.Allow(key)
	if allowed {
		t.Fatal("expected request to be rate limited")
	}

	// Test 3: Wait for refill
	time.Sleep(1 * time.Second)
	
	// Should have ~10 tokens now
	allowedCount := 0
	for i := 0; i < 15; i++ {
		if allowed, _ := tb.Allow(key); allowed {
			allowedCount++
		}
	}
	
	if allowedCount < 8 || allowedCount > 12 {
		t.Fatalf("expected ~10 requests, got %d", allowedCount)
	}
}

func TestSlidingWindowLog(t *testing.T) {
	// 5 requests per second
	swl := NewSlidingWindowLog(5, time.Second)
	key := "test-api"

	// Test 1: Allow 5 requests
	for i := 0; i < 5; i++ {
		allowed, err := swl.Allow(key)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if !allowed {
			t.Fatalf("expected request %d to be allowed", i+1)
		}
	}

	// Test 2: 6th request should fail
	allowed, _ := swl.Allow(key)
	if allowed {
		t.Fatal("expected 6th request to be rate limited")
	}

	// Test 3: Wait for window to slide
	time.Sleep(1100 * time.Millisecond)
	
	// Should allow new requests
	allowed, _ = swl.Allow(key)
	if !allowed {
		t.Fatal("expected request after window slide to be allowed")
	}
}

func TestConcurrentAccess(t *testing.T) {
	tb := NewTokenBucket(100, 200)
	
	var allowed int32
	var rejected int32
	var wg sync.WaitGroup
	
	// Simulate 10 concurrent clients
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(clientID int) {
			defer wg.Done()
			key := fmt.Sprintf("client-%d", clientID)
			
			for j := 0; j < 50; j++ {
				if ok, _ := tb.Allow(key); ok {
					atomic.AddInt32(&allowed, 1)
				} else {
					atomic.AddInt32(&rejected, 1)
				}
				time.Sleep(10 * time.Millisecond)
			}
		}(i)
	}
	
	wg.Wait()
	
	total := allowed + rejected
	if total != 500 {
		t.Fatalf("expected 500 total requests, got %d", total)
	}
	
	t.Logf("Concurrent test: %d allowed, %d rejected", allowed, rejected)
}

func TestMultipleKeys(t *testing.T) {
	tb := NewTokenBucket(5, 10)
	
	// Different users should have independent limits
	users := []string{"alice", "bob", "charlie"}
	
	for _, user := range users {
		// Each user should get their full burst
		for i := 0; i < 10; i++ {
			allowed, _ := tb.Allow(user)
			if !allowed {
				t.Fatalf("expected request %d for user %s to be allowed", i+1, user)
			}
		}
		
		// 11th request should fail
		allowed, _ := tb.Allow(user)
		if allowed {
			t.Fatalf("expected 11th request for user %s to be rejected", user)
		}
	}
}

func BenchmarkTokenBucket(b *testing.B) {
	tb := NewTokenBucket(1000, 2000)
	key := "bench-user"
	
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			tb.Allow(key)
		}
	})
}

func BenchmarkSlidingWindowLog(b *testing.B) {
	swl := NewSlidingWindowLog(1000, time.Second)
	key := "bench-user"
	
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			swl.Allow(key)
		}
	})
}

// Example of testing with simulated time
type MockClock struct {
	now time.Time
	mu  sync.Mutex
}

func (m *MockClock) Now() time.Time {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.now
}

func (m *MockClock) Advance(d time.Duration) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.now = m.now.Add(d)
}