stock-bot/libs/core/queue/test/rate-limiter.test.ts

import { afterEach, beforeEach, describe, expect, test } from 'bun:test';
import Redis from 'ioredis';
import { QueueRateLimiter } from '../src/rate-limiter';
import { getRedisConnection } from '../src/utils';

// Suppress Redis connection errors in tests
process.on('unhandledRejection', (reason, promise) => {
  if (reason && typeof reason === 'object' && 'message' in reason) {
    const message = (reason as Error).message;
    if (
      message.includes('Connection is closed') ||
      message.includes('Connection is in monitoring mode')
    ) {
      return;
    }
  }
  console.error('Unhandled Rejection at:', promise, 'reason:', reason);
});

describe('QueueRateLimiter', () => {
  let redisClient: Redis;
  let rateLimiter: QueueRateLimiter;

  const redisConfig = {
    host: 'localhost',
    port: 6379,
    password: '',
    db: 0,
  };

  beforeEach(async () => {
    // Create Redis client
    redisClient = new Redis(getRedisConnection(redisConfig));

    // Clear Redis keys for tests
    try {
      const keys = await redisClient.keys('rl:*');
      if (keys.length > 0) {
        await redisClient.del(...keys);
      }
    } catch {
      // Ignore cleanup errors
    }
    rateLimiter = new QueueRateLimiter(redisClient);
  });

  afterEach(async () => {
    if (redisClient) {
      try {
        await redisClient.quit();
      } catch {
        // Ignore cleanup errors
      }
    }
    await new Promise(resolve => setTimeout(resolve, 50));
  });

  describe('Rate Limit Rules', () => {
    test('should add and enforce global rate limit', async () => {
      rateLimiter.addRule({
        level: 'global',
        config: {
          points: 5,
          duration: 1, // 1 second
        },
      });

      // Consume 5 points
      for (let i = 0; i < 5; i++) {
        const result = await rateLimiter.checkLimit('any-handler', 'any-operation');
        expect(result.allowed).toBe(true);
      }

      // 6th request should be blocked
      const blocked = await rateLimiter.checkLimit('any-handler', 'any-operation');
      expect(blocked.allowed).toBe(false);
      expect(blocked.retryAfter).toBeGreaterThan(0);
    });

    test('should add and enforce handler-level rate limit', async () => {
      rateLimiter.addRule({
        level: 'handler',
        handler: 'api-handler',
        config: {
          points: 3,
          duration: 1,
        },
      });

      // api-handler should be limited
      for (let i = 0; i < 3; i++) {
        const result = await rateLimiter.checkLimit('api-handler', 'any-operation');
        expect(result.allowed).toBe(true);
      }

      const blocked = await rateLimiter.checkLimit('api-handler', 'any-operation');
      expect(blocked.allowed).toBe(false);

      // Other handlers should not be limited
      const otherHandler = await rateLimiter.checkLimit('other-handler', 'any-operation');
      expect(otherHandler.allowed).toBe(true);
    });

    test('should add and enforce operation-level rate limit', async () => {
      rateLimiter.addRule({
        level: 'operation',
        handler: 'data-handler',
        operation: 'fetch-prices',
        config: {
          points: 2,
          duration: 1,
        },
      });

      // Specific operation should be limited
      for (let i = 0; i < 2; i++) {
        const result = await rateLimiter.checkLimit('data-handler', 'fetch-prices');
        expect(result.allowed).toBe(true);
      }

      const blocked = await rateLimiter.checkLimit('data-handler', 'fetch-prices');
      expect(blocked.allowed).toBe(false);

      // Other operations on same handler should work
      const otherOp = await rateLimiter.checkLimit('data-handler', 'fetch-volume');
      expect(otherOp.allowed).toBe(true);
    });

    test('should enforce multiple rate limits (most restrictive wins)', async () => {
      // Global: 10/sec
      rateLimiter.addRule({
        level: 'global',
        config: { points: 10, duration: 1 },
      });

      // Handler: 5/sec
      rateLimiter.addRule({
        level: 'handler',
        handler: 'test-handler',
        config: { points: 5, duration: 1 },
      });

      // Operation: 2/sec
      rateLimiter.addRule({
        level: 'operation',
        handler: 'test-handler',
        operation: 'test-op',
        config: { points: 2, duration: 1 },
      });

      // Should be limited by operation level (most restrictive)
      for (let i = 0; i < 2; i++) {
        const result = await rateLimiter.checkLimit('test-handler', 'test-op');
        expect(result.allowed).toBe(true);
      }

      const blocked = await rateLimiter.checkLimit('test-handler', 'test-op');
      expect(blocked.allowed).toBe(false);
    });
  });

  describe('Rate Limit Status', () => {
    test('should get rate limit status', async () => {
      rateLimiter.addRule({
        level: 'handler',
        handler: 'status-test',
        config: { points: 10, duration: 60 },
      });

      // Consume some points
      await rateLimiter.checkLimit('status-test', 'operation');
      await rateLimiter.checkLimit('status-test', 'operation');

      const status = await rateLimiter.getStatus('status-test', 'operation');
      expect(status.handler).toBe('status-test');
      expect(status.operation).toBe('operation');
      expect(status.limits.length).toBe(1);
      expect(status.limits[0].points).toBe(10);
      expect(status.limits[0].remaining).toBe(8);
    });

    test('should show multiple applicable limits in status', async () => {
      rateLimiter.addRule({
        level: 'global',
        config: { points: 100, duration: 60 },
      });

      rateLimiter.addRule({
        level: 'handler',
        handler: 'multi-test',
        config: { points: 50, duration: 60 },
      });

      const status = await rateLimiter.getStatus('multi-test', 'operation');
      expect(status.limits.length).toBe(2);

      const globalLimit = status.limits.find(l => l.level === 'global');
      const handlerLimit = status.limits.find(l => l.level === 'handler');

      expect(globalLimit?.points).toBe(100);
      expect(handlerLimit?.points).toBe(50);
    });
  });

  describe('Rate Limit Management', () => {
    test('should reset rate limits', async () => {
      rateLimiter.addRule({
        level: 'handler',
        handler: 'reset-test',
        config: { points: 1, duration: 60 },
      });

      // Consume the limit
      await rateLimiter.checkLimit('reset-test', 'operation');
      const blocked = await rateLimiter.checkLimit('reset-test', 'operation');
      expect(blocked.allowed).toBe(false);

      // Reset limits
      await rateLimiter.reset('reset-test');

      // Should be allowed again
      const afterReset = await rateLimiter.checkLimit('reset-test', 'operation');
      expect(afterReset.allowed).toBe(true);
    });

    test('should get all rules', async () => {
      rateLimiter.addRule({
        level: 'global',
        config: { points: 100, duration: 60 },
      });

      rateLimiter.addRule({
        level: 'handler',
        handler: 'test',
        config: { points: 50, duration: 60 },
      });

      const rules = rateLimiter.getRules();
      expect(rules.length).toBe(2);
      expect(rules[0].level).toBe('global');
      expect(rules[1].level).toBe('handler');
    });

    test('should remove specific rule', async () => {
      rateLimiter.addRule({
        level: 'handler',
        handler: 'remove-test',
        config: { points: 1, duration: 1 },
      });

      // Verify rule exists
      await rateLimiter.checkLimit('remove-test', 'op');
      const blocked = await rateLimiter.checkLimit('remove-test', 'op');
      expect(blocked.allowed).toBe(false);

      // Remove rule
      const removed = rateLimiter.removeRule('handler', 'remove-test');
      expect(removed).toBe(true);

      // Should not be limited anymore
      const afterRemove = await rateLimiter.checkLimit('remove-test', 'op');
      expect(afterRemove.allowed).toBe(true);
    });
  });

  describe('Block Duration', () => {
    test('should block for specified duration after limit exceeded', async () => {
      rateLimiter.addRule({
        level: 'handler',
        handler: 'block-test',
        config: {
          points: 1,
          duration: 1,
          blockDuration: 2, // Block for 2 seconds
        },
      });

      // Consume limit
      await rateLimiter.checkLimit('block-test', 'op');

      // Should be blocked
      const blocked = await rateLimiter.checkLimit('block-test', 'op');
      expect(blocked.allowed).toBe(false);
      expect(blocked.retryAfter).toBeGreaterThanOrEqual(1000); // At least 1 second
    });
  });

  describe('Error Handling', () => {
    test('should allow requests when rate limiter fails', async () => {
      // Create a rate limiter with invalid redis client
      const badRedis = new Redis({
        host: 'invalid-host',
        port: 9999,
        retryStrategy: () => null, // Disable retries
      });

      const failingLimiter = new QueueRateLimiter(badRedis);

      failingLimiter.addRule({
        level: 'global',
        config: { points: 1, duration: 1 },
      });

      // Should allow even though Redis is not available
      const result = await failingLimiter.checkLimit('test', 'test');
      expect(result.allowed).toBe(true);

      badRedis.disconnect();
    });
  });
});