agents_and_robots/tools/ratelimit_test.go

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