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