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

func TestRateLimiterWaitDoesNotConsumeGlobalCapacityWhileChatIsBlocked(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()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()
	if err := rl.Wait(ctx, 42); err == nil {
		t.Fatal("expected blocked chat wait to fail")
	}
	if !rl.GlobalAllow() {
		t.Fatal("blocked chat wait consumed global capacity")
	}
}

func TestRateLimiterWaitObservesExtendedCooldowns(t *testing.T) {
	tests := []struct {
		name string
		wait func(*RateLimiter, context.Context) error
		set  func(*RateLimiter, time.Time)
	}{
		{
			name: "global",
			wait: func(rl *RateLimiter, ctx context.Context) error {
				return rl.waitForGlobalUnlock(ctx)
			},
			set: func(rl *RateLimiter, until time.Time) {
				rl.globalMu.Lock()
				rl.globalLockUntil = until
				rl.globalMu.Unlock()
			},
		},
		{
			name: "chat",
			wait: func(rl *RateLimiter, ctx context.Context) error {
				return rl.waitForChatUnlock(ctx, 42)
			},
			set: func(rl *RateLimiter, until time.Time) {
				rl.chatMu.Lock()
				rl.chatLocks[42] = until
				rl.chatMu.Unlock()
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			rl := NewRateLimiter()
			tt.set(rl, time.Now().Add(20*time.Millisecond))
			started := time.Now()
			done := make(chan error, 1)
			go func() { done <- tt.wait(rl, context.Background()) }()

			time.Sleep(5 * time.Millisecond)
			tt.set(rl, time.Now().Add(70*time.Millisecond))
			if err := <-done; err != nil {
				t.Fatalf("wait returned error: %v", err)
			}
			if elapsed := time.Since(started); elapsed < 60*time.Millisecond {
				t.Fatalf("wait ignored extended cooldown: %v", elapsed)
			}
		})
	}
}

func TestRateLimiterCleanupIgnoresNonPositiveThreshold(t *testing.T) {
	rl := NewRateLimiter()
	want := rl.getChatLimiter(42)
	rl.Cleanup(0)
	if got := rl.getChatLimiter(42); got != want {
		t.Fatal("Cleanup(0) replaced an active limiter")
	}
}

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

func TestRateLimiterCleanupPreservesAcquiredLimiter(t *testing.T) {
	rl := NewRateLimiter()
	limiter, release := rl.acquireChatLimiter(42)

	rl.chatMu.Lock()
	rl.chatLastSeen[42] = time.Now().Add(-time.Hour)
	rl.chatMu.Unlock()
	rl.Cleanup(time.Minute)

	if got := rl.getChatLimiter(42); got != limiter {
		t.Fatal("Cleanup replaced a limiter while it was acquired")
	}
	release()
}