RedditBot/tests/rate-limiter.test.ts

import { describe, it, expect, beforeEach, vi } from 'vitest';
import { RateLimiter, type SentMessageHistoryProvider } from '../src/utils/rate-limiter.js';
import { createMockLogger } from '../tests/mockLogger.js';
import { getTestConfig } from '../tests/testUtils.js';

describe('RateLimiter', () => {
  let limiter: RateLimiter;
  let logger: ReturnType<typeof createMockLogger>;

  beforeEach(() => {
    logger = createMockLogger();
    const config = getTestConfig({
      rateLimit: { minDelayBetweenMessagesMs: 100, maxMessagesPerHour: 10, jitterMs: 50 },
    });
    limiter = new RateLimiter(config, logger);
  });

  it('allows immediate first call', async () => {
    const start = Date.now();
    await limiter.throttle(() => Promise.resolve('done'));
    const elapsed = Date.now() - start;
    expect(elapsed).toBeLessThan(50);
  });

  it('enforces minimum delay between messages', async () => {
    await limiter.throttle(() => Promise.resolve('first'));

    const start = Date.now();
    await limiter.throttle(() => Promise.resolve('second'));
    const elapsed = Date.now() - start;

    expect(elapsed).toBeGreaterThanOrEqual(100);
    expect(elapsed).toBeLessThan(200); // 100ms delay + 50ms max jitter
  });

  it('adds jitter to the delay', async () => {
    // Run a message first to set lastMessageTime
    await limiter.throttle(() => Promise.resolve('first'));

    const delays: number[] = [];
    for (let i = 0; i < 5; i++) {
      const start = Date.now();
      await limiter.throttle(() => Promise.resolve(`msg-${i}`));
      delays.push(Date.now() - start);
    }

    // At least one delay should differ significantly from the base delay
    const uniqueDelays = new Set(delays.map(d => Math.floor(d / 10)));
    expect(uniqueDelays.size).toBeGreaterThan(1);
  });

  it('resets hourly counter', async () => {
    const config = getTestConfig({
      rateLimit: { minDelayBetweenMessagesMs: 10, maxMessagesPerHour: 3, jitterMs: 10 },
    });
    // 500ms "hour" so we can actually observe the reset wait inside the test budget.
    limiter = new RateLimiter(config, logger, 500);

    for (let i = 0; i < 3; i++) {
      await limiter.throttle(() => Promise.resolve(`msg-${i}`));
    }

    // 4th should be blocked by hourly limit and wait for the reset.
    const start = Date.now();
    await limiter.throttle(() => Promise.resolve('blocked'));
    const elapsed = Date.now() - start;

    expect(elapsed).toBeGreaterThanOrEqual(350);
  });

  it('counts only when shouldCount returns true', async () => {
    const config = getTestConfig({
      rateLimit: { minDelayBetweenMessagesMs: 10, maxMessagesPerHour: 2, jitterMs: 10 },
    });
    limiter = new RateLimiter(config, logger, 500);

    // These should NOT be counted
    await limiter.throttle(() => Promise.resolve('skip'), () => false);
    await limiter.throttle(() => Promise.resolve('skip'), () => false);

    // These SHOULD be counted
    await limiter.throttle(() => Promise.resolve('count'), () => true);
    await limiter.throttle(() => Promise.resolve('count'), () => true);

    // Next counted one should be delayed
    const start = Date.now();
    await limiter.throttle(() => Promise.resolve('blocked'), () => true);
    const elapsed = Date.now() - start;
    expect(elapsed).toBeGreaterThanOrEqual(10);
  });

  it('returns the function result', async () => {
    const result = await limiter.throttle(() => Promise.resolve({ value: 42 }));
    expect(result).toEqual({ value: 42 });
  });

  it('constrains to concurrency of 1', async () => {
    const order: number[] = [];
    const promises = [
      limiter.throttle(() => { order.push(1); return Promise.resolve('a'); }),
      limiter.throttle(() => { order.push(2); return Promise.resolve('b'); }),
      limiter.throttle(() => { order.push(3); return Promise.resolve('c'); }),
    ];
    await Promise.all(promises);
    expect(order).toEqual([1, 2, 3]);
  });

  describe('getStats', () => {
    it('returns initial stats', () => {
      const stats = limiter.getStats();
      expect(stats.messagesThisHour).toBe(0);
      // minutesSinceLastMessage should be large initially
      expect(stats.minutesSinceLastMessage).toBeGreaterThan(0);
    });

    it('updates stats after messages', async () => {
      await limiter.throttle(() => Promise.resolve('msg'));
      const stats = limiter.getStats();
      expect(stats.messagesThisHour).toBe(1);
      expect(stats.minutesSinceLastMessage).toBeLessThan(1);
    });
  });

  describe('persisted send history', () => {
    it('seeds the rolling hourly count from successful sends persisted before restart', async () => {
      const now = Date.now();
      const history: SentMessageHistoryProvider = vi.fn().mockResolvedValue([
        now - 5_000,
        now - 10_000,
        now - 3_700_000,
      ]);
      const config = getTestConfig({
        rateLimit: { minDelayBetweenMessagesMs: 0, maxMessagesPerHour: 3, jitterMs: 0 },
      });
      limiter = new RateLimiter(config, logger, 3_600_000, { sentMessageHistory: history });

      await limiter.initialize();

      expect(history).toHaveBeenCalledWith(expect.any(Number));
      expect(limiter.getStats().messagesThisHour).toBe(2);
      expect(limiter.getTimers().messagesRemaining).toBe(1);
    });

    it('initializes persisted history only once when concurrent sends arrive', async () => {
      const history: SentMessageHistoryProvider = vi.fn().mockResolvedValue([]);
      limiter = new RateLimiter(getTestConfig(), logger, 3_600_000, { sentMessageHistory: history });

      await Promise.all([
        limiter.throttle(() => Promise.resolve('one')),
        limiter.throttle(() => Promise.resolve('two')),
      ]);

      expect(history).toHaveBeenCalledTimes(1);
    });

    it('cancels an hourly-limit wait when shutdown aborts it', async () => {
      const controller = new AbortController();
      const limiterWithHistory = new RateLimiter(
        getTestConfig({ rateLimit: { minDelayBetweenMessagesMs: 0, maxMessagesPerHour: 1, jitterMs: 0 } }),
        logger,
        60_000,
        { sentMessageHistory: async () => [Date.now()] },
      );
      await limiterWithHistory.initialize();
      const pending = limiterWithHistory.throttle(() => Promise.resolve('should not run'), () => true, controller.signal);
      controller.abort();

      await expect(pending).rejects.toMatchObject({ name: 'AbortError' });
    });
  });

  describe('getTimers', () => {
    it('returns null nextMessageAvailableAt initially', () => {
      const timers = limiter.getTimers();
      expect(timers.nextMessageAvailableAt).toBeNull();
      expect(timers.isHourlyLimitReached).toBe(false);
      expect(timers.messagesRemaining).toBe(10);
    });

    it('returns nextMessageAvailableAt after a message', async () => {
      await limiter.throttle(() => Promise.resolve('msg'));
      const timers = limiter.getTimers();
      expect(timers.nextMessageAvailableAt).toBeDefined();
      expect(timers.nextMessageAvailableAt!).toBeGreaterThan(Date.now());
    });

    it('shows hourly limit reached when max messages sent', async () => {
      const config = getTestConfig({
        rateLimit: { minDelayBetweenMessagesMs: 10, maxMessagesPerHour: 2, jitterMs: 10 },
      });
      limiter = new RateLimiter(config, logger);

      await limiter.throttle(() => Promise.resolve('msg1'));
      await limiter.throttle(() => Promise.resolve('msg2'));

      const timers = limiter.getTimers();
      expect(timers.isHourlyLimitReached).toBe(true);
      expect(timers.messagesRemaining).toBe(0);
    });

    it('calculates messages remaining', async () => {
      const config = getTestConfig({
        rateLimit: { minDelayBetweenMessagesMs: 10, maxMessagesPerHour: 5, jitterMs: 10 },
      });
      limiter = new RateLimiter(config, logger);

      await limiter.throttle(() => Promise.resolve('msg1'));
      await limiter.throttle(() => Promise.resolve('msg2'));

      const timers = limiter.getTimers();
      expect(timers.messagesRemaining).toBe(3);
    });
  });

  describe('edge cases', () => {
    it('handles function that throws', async () => {
      await expect(limiter.throttle(() => { throw new Error('fail'); })).rejects.toThrow('fail');
    });

    it('handles async function errors', async () => {
      await expect(limiter.throttle(() => Promise.reject(new Error('async fail')))).rejects.toThrow('async fail');
    });

    it('handles synchronous return values', async () => {
      const result = await limiter.throttle(() => 'sync value');
      expect(result).toBe('sync value');
    });

    it('handles zero minDelay', async () => {
      const config = getTestConfig({
        rateLimit: { minDelayBetweenMessagesMs: 0, maxMessagesPerHour: 100, jitterMs: 0 },
      });
      limiter = new RateLimiter(config, logger);

      const start = Date.now();
      for (let i = 0; i < 5; i++) {
        await limiter.throttle(() => Promise.resolve(`msg-${i}`));
      }
      const elapsed = Date.now() - start;
      expect(elapsed).toBeLessThan(100);
    });
  });
});