diff --git a/tools/ratelimit_test.go b/tools/ratelimit_test.go new file mode 100644 index 0000000..3417afb --- /dev/null +++ b/tools/ratelimit_test.go @@ -0,0 +1,161 @@ +package tools + +import ( + "context" + "log/slog" + "testing" + "time" +) + +func TestRateLimiter_AllowWithinLimit(t *testing.T) { + rl := NewRateLimiter(3, time.Minute) + + for i := 0; i < 3; i++ { + if !rl.Allow("room1") { + t.Fatalf("call %d should be allowed", i+1) + } + } +} + +func TestRateLimiter_DenyOverLimit(t *testing.T) { + rl := NewRateLimiter(3, time.Minute) + + for i := 0; i < 3; i++ { + rl.Allow("room1") + } + + if rl.Allow("room1") { + t.Fatal("4th call should be denied") + } +} + +func TestRateLimiter_DifferentKeysIndependent(t *testing.T) { + rl := NewRateLimiter(2, time.Minute) + + rl.Allow("room1") + rl.Allow("room1") + + // room1 is full, but room2 should still be allowed + if rl.Allow("room1") { + t.Fatal("room1 3rd call should be denied") + } + if !rl.Allow("room2") { + t.Fatal("room2 should be allowed independently") + } +} + +func TestRateLimiter_WindowExpiry(t *testing.T) { + // Use a very short window for testing + rl := NewRateLimiter(2, 50*time.Millisecond) + + rl.Allow("room1") + rl.Allow("room1") + + if rl.Allow("room1") { + t.Fatal("should be denied before window expires") + } + + // Wait for window to expire + time.Sleep(60 * time.Millisecond) + + if !rl.Allow("room1") { + t.Fatal("should be allowed after window expires") + } +} + +func TestRateLimiter_Cleanup(t *testing.T) { + rl := NewRateLimiter(5, 50*time.Millisecond) + + rl.Allow("room1") + rl.Allow("room2") + + // Wait for entries to expire + time.Sleep(60 * time.Millisecond) + + rl.Cleanup() + + rl.mu.Lock() + count := len(rl.buckets) + rl.mu.Unlock() + + if count != 0 { + t.Fatalf("expected 0 buckets after cleanup, got %d", count) + } +} + +func TestRateLimiter_CleanupKeepsActive(t *testing.T) { + rl := NewRateLimiter(5, time.Minute) + + rl.Allow("room1") + + rl.Cleanup() + + rl.mu.Lock() + count := len(rl.buckets) + rl.mu.Unlock() + + if count != 1 { + t.Fatalf("expected 1 bucket after cleanup of active entries, got %d", count) + } +} + +func TestRegistry_ExecuteForRoom_RateLimited(t *testing.T) { + logger := slog.Default() + reg := NewRegistry(logger) + + // Register a simple echo tool + reg.Register(Tool{ + Def: Def{Name: "echo", Description: "echo tool"}, + Exec: func(_ context.Context, args map[string]any) Result { + return Result{Output: "ok"} + }, + }) + + rl := NewRateLimiter(2, time.Minute) + reg.SetRateLimiter(rl) + + ctx := context.Background() + + // First two calls succeed + r1 := reg.ExecuteForRoom(ctx, "echo", "", "!room:test") + if r1.Err != nil { + t.Fatalf("call 1 should succeed: %v", r1.Err) + } + r2 := reg.ExecuteForRoom(ctx, "echo", "", "!room:test") + if r2.Err != nil { + t.Fatalf("call 2 should succeed: %v", r2.Err) + } + + // Third call is rate limited + r3 := reg.ExecuteForRoom(ctx, "echo", "", "!room:test") + if r3.Err == nil { + t.Fatal("call 3 should be rate limited") + } + + // Different room still works + r4 := reg.ExecuteForRoom(ctx, "echo", "", "!room:other") + if r4.Err != nil { + t.Fatalf("different room should succeed: %v", r4.Err) + } +} + +func TestRegistry_ExecuteForRoom_NoLimiter(t *testing.T) { + logger := slog.Default() + reg := NewRegistry(logger) + + reg.Register(Tool{ + Def: Def{Name: "echo", Description: "echo tool"}, + Exec: func(_ context.Context, args map[string]any) Result { + return Result{Output: "ok"} + }, + }) + + // No rate limiter set — all calls should succeed + ctx := context.Background() + for i := 0; i < 20; i++ { + r := reg.ExecuteForRoom(ctx, "echo", "", "!room:test") + if r.Err != nil { + t.Fatalf("call %d should succeed without limiter: %v", i+1, r.Err) + } + } +}