Laniakea/utils/limiter_test.go

package utils

import (
	"context"
	"errors"
	"testing"
	"time"

	"golang.org/x/time/rate"
)

func TestRateLimiterCheckDropOverflowHonorsGlobalLock(t *testing.T) {
	rl := NewRateLimiter()
	rl.SetGlobalLock(1)

	if err := rl.Check(context.Background(), true, 0); !errors.Is(err, ErrDropOverflow) {
		t.Fatalf("expected ErrDropOverflow, got %v", err)
	}
}

func TestRateLimiterChatLocksAreScopedPerChat(t *testing.T) {
	rl := NewRateLimiter()
	rl.SetChatLock(42, 1)

	if rl.Allow(42) {
		t.Fatal("expected locked chat to be rejected")
	}
	if !rl.Allow(7) {
		t.Fatal("expected unrelated chat to remain allowed")
	}
}

func TestRateLimiterRejectedChatDoesNotConsumeGlobalCapacity(t *testing.T) {
	rl := NewRateLimiter()
	rl.SetGlobalRate(1)

	if !rl.Allow(42) {
		t.Fatal("expected initial request for chat 42 to succeed")
	}
	rl.globalMu.Lock()
	rl.globalLimiter = rate.NewLimiter(1, 1)
	rl.globalMu.Unlock()

	if rl.Allow(42) {
		t.Fatal("expected exhausted chat limiter to reject the request")
	}
	if !rl.Allow(7) {
		t.Fatal("expected rejected chat request not to consume global capacity")
	}
}

func TestRateLimiterGlobalWaitRespectsContextCancellation(t *testing.T) {
	rl := NewRateLimiter()
	rl.SetGlobalLock(1)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()

	if err := rl.GlobalWait(ctx); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("expected DeadlineExceeded, got %v", err)
	}
}

// TestRateLimiterCleanupEvictsIdleChats guards the memory-leak fix: per-chat
// limiter and lastSeen state must be reclaimed by Cleanup once the entry has
// been idle for longer than the threshold, while still-active chats and
// unexpired cooldowns must survive.
func TestRateLimiterCleanupEvictsIdleChats(t *testing.T) {
	rl := NewRateLimiter()

	// Touch chat 1 to make it tracked, then backdate its last-seen marker
	// so it looks idle from Cleanup's perspective.
	if !rl.Allow(1) {
		t.Fatal("expected initial Allow for chat 1 to succeed")
	}
	rl.chatMu.Lock()
	rl.chatLastSeen[1] = time.Now().Add(-time.Hour)
	rl.chatMu.Unlock()

	// Touch chat 2 so it stays "active".
	if !rl.Allow(2) {
		t.Fatal("expected initial Allow for chat 2 to succeed")
	}

	// Expired cooldown should be evicted; future cooldown should survive.
	rl.SetChatLock(10, 1)
	rl.chatMu.Lock()
	rl.chatLocks[10] = time.Now().Add(-time.Second)
	rl.chatLocks[11] = time.Now().Add(time.Hour)
	rl.chatMu.Unlock()

	rl.Cleanup(time.Minute)

	rl.chatMu.RLock()
	defer rl.chatMu.RUnlock()
	if _, ok := rl.chatLimiters[1]; ok {
		t.Fatal("expected idle chat 1 limiter to be evicted")
	}
	if _, ok := rl.chatLastSeen[1]; ok {
		t.Fatal("expected idle chat 1 lastSeen to be evicted")
	}
	if _, ok := rl.chatLimiters[2]; !ok {
		t.Fatal("expected active chat 2 limiter to remain")
	}
	if _, ok := rl.chatLocks[10]; ok {
		t.Fatal("expected expired chat 10 lock to be evicted")
	}
	if _, ok := rl.chatLocks[11]; !ok {
		t.Fatal("expected future chat 11 lock to remain")
	}
}