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