proffesor-for-testing__agentic-qe/tests/unit/mcp/security/rate-limiter.test.ts

/**
 * Agentic QE v3 - Rate Limiter Tests
 * Tests for token bucket rate limiting
 */

import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import {
  RateLimiter,
  SlidingWindowRateLimiter,
  createRateLimiter,
  createStrictRateLimiter,
  createSlidingWindowLimiter,
} from '../../../../src/mcp/security/rate-limiter';

describe('RateLimiter', () => {
  let limiter: RateLimiter;

  beforeEach(() => {
    limiter = createRateLimiter({
      tokensPerSecond: 10,
      maxBurst: 20,
      perClient: true,
    });
  });

  afterEach(() => {
    limiter.dispose();
  });

  describe('basic rate limiting', () => {
    it('should allow requests within limit', () => {
      const result = limiter.check('client1');

      expect(result.allowed).toBe(true);
      expect(result.remaining).toBeLessThanOrEqual(20);
    });

    it('should track remaining tokens', () => {
      // Use some tokens
      limiter.check('client1');
      limiter.check('client1');
      limiter.check('client1');

      const result = limiter.check('client1');
      expect(result.remaining).toBe(16); // 20 - 4 = 16
    });

    it('should deny requests when tokens exhausted', () => {
      // Exhaust all tokens
      for (let i = 0; i < 20; i++) {
        limiter.check('client1');
      }

      const result = limiter.check('client1');
      expect(result.allowed).toBe(false);
      expect(result.remaining).toBe(0);
      expect(result.retryAfter).toBeDefined();
    });

    it('should provide retry-after time', () => {
      // Exhaust tokens
      for (let i = 0; i < 20; i++) {
        limiter.check('client1');
      }

      const result = limiter.check('client1');
      expect(result.retryAfter).toBeGreaterThan(0);
      expect(result.retryAfter).toBeLessThanOrEqual(1000); // Max 1 second wait for 10 tokens/sec
    });
  });

  describe('per-client limiting', () => {
    it('should maintain separate buckets per client', () => {
      // Exhaust client1's tokens
      for (let i = 0; i < 20; i++) {
        limiter.check('client1');
      }

      // client2 should still have tokens
      const result = limiter.check('client2');
      expect(result.allowed).toBe(true);
      expect(result.remaining).toBe(19);
    });

    it('should use global bucket when no client specified', () => {
      // Create limiter without per-client limiting
      const globalLimiter = createRateLimiter({
        tokensPerSecond: 10,
        maxBurst: 20,
        perClient: false,
      });

      const result1 = globalLimiter.check();
      const result2 = globalLimiter.check();

      expect(result1.remaining - result2.remaining).toBe(1);

      globalLimiter.dispose();
    });
  });

  describe('token refill', () => {
    it('should refill tokens over time', async () => {
      vi.useFakeTimers();

      // Use some tokens
      for (let i = 0; i < 10; i++) {
        limiter.check('client1');
      }

      const beforeRefill = limiter.check('client1');
      expect(beforeRefill.remaining).toBe(9);

      // Wait for tokens to refill (100ms = 1 token at 10/sec)
      vi.advanceTimersByTime(1000);

      const afterRefill = limiter.check('client1');
      expect(afterRefill.remaining).toBeGreaterThan(9);

      vi.useRealTimers();
    });

    it('should not exceed max burst', async () => {
      vi.useFakeTimers();

      // Wait a long time - should still cap at maxBurst
      vi.advanceTimersByTime(10000);

      const result = limiter.check('client1');
      expect(result.remaining).toBe(19); // maxBurst - 1
      expect(result.remaining).toBeLessThanOrEqual(20);

      vi.useRealTimers();
    });
  });

  describe('consume multiple tokens', () => {
    it('should consume multiple tokens at once', () => {
      const result = limiter.consume('client1', undefined, 5);

      expect(result.allowed).toBe(true);
      expect(result.remaining).toBe(15);
    });

    it('should deny if not enough tokens', () => {
      const result = limiter.consume('client1', undefined, 25);

      expect(result.allowed).toBe(false);
    });
  });

  describe('async wait', () => {
    it('should wait for tokens to become available', async () => {
      vi.useFakeTimers();

      // Exhaust tokens
      for (let i = 0; i < 20; i++) {
        limiter.check('client1');
      }

      const waitPromise = limiter.wait('client1');

      // Advance time to allow token refill
      vi.advanceTimersByTime(200);

      const result = await waitPromise;
      expect(result.allowed).toBe(true);

      vi.useRealTimers();
    });
  });

  describe('endpoint-specific limits', () => {
    it('should apply endpoint-specific limits', () => {
      limiter.addEndpointLimit({
        pattern: /\/api\/admin/,
        tokensPerSecond: 5,
        maxBurst: 10,
      });

      const adminResult = limiter.check('client1', '/api/admin/users');
      // Token bucket may have slight timing variations, use range check
      expect(adminResult.remaining).toBeGreaterThanOrEqual(8); // ~10 maxBurst - 1 (±1 for timing)
      expect(adminResult.remaining).toBeLessThanOrEqual(10);

      // Normal endpoint has separate bucket from admin endpoint
      const normalResult = limiter.check('client1', '/api/users');
      expect(normalResult.remaining).toBeGreaterThanOrEqual(18); // ~20 maxBurst - 1 (±1 for timing)
      expect(normalResult.remaining).toBeLessThanOrEqual(20);
    });

    it('should clear endpoint limits', () => {
      limiter.addEndpointLimit({
        pattern: '/test',
        tokensPerSecond: 1,
        maxBurst: 2,
      });

      limiter.clearEndpointLimits();

      const result = limiter.check('client1', '/test');
      expect(result.remaining).toBe(19); // Uses default maxBurst
    });
  });

  describe('reset', () => {
    it('should reset client bucket', () => {
      // Use tokens
      for (let i = 0; i < 10; i++) {
        limiter.check('client1');
      }

      limiter.resetClient('client1');

      const result = limiter.check('client1');
      expect(result.remaining).toBe(19); // Full bucket
    });

    it('should reset all buckets', () => {
      limiter.check('client1');
      limiter.check('client2');
      limiter.check('client3');

      limiter.reset();

      const stats = limiter.getStats();
      expect(stats.activeBuckets).toBe(1); // Only global bucket
      expect(stats.totalRequests).toBe(0);
    });
  });

  describe('statistics', () => {
    it('should track request statistics', () => {
      limiter.check('client1');
      limiter.check('client2');
      limiter.check('client1');

      const stats = limiter.getStats();

      expect(stats.totalRequests).toBe(3);
      expect(stats.allowedRequests).toBe(3);
      expect(stats.deniedRequests).toBe(0);
    });

    it('should track denied requests', () => {
      // Exhaust tokens
      for (let i = 0; i < 20; i++) {
        limiter.check('client1');
      }

      // Try more requests
      limiter.check('client1');
      limiter.check('client1');

      const stats = limiter.getStats();
      expect(stats.deniedRequests).toBe(2);
    });

    it('should track active buckets', () => {
      limiter.check('client1');
      limiter.check('client2');
      limiter.check('client3');

      const stats = limiter.getStats();
      expect(stats.activeBuckets).toBe(4); // 3 clients + 1 global
    });

    it('should get client-specific stats', () => {
      limiter.check('client1');
      limiter.check('client1');
      limiter.check('client1');

      const clientStats = limiter.getClientStats('client1');

      expect(clientStats).not.toBeNull();
      expect(clientStats?.totalRequests).toBe(3);
    });

    it('should return null for unknown client', () => {
      const stats = limiter.getClientStats('unknown');
      expect(stats).toBeNull();
    });
  });

  describe('rate limit headers', () => {
    it('should provide rate limit headers', () => {
      const result = limiter.check('client1');

      expect(result.headers['X-RateLimit-Limit']).toBe(20);
      expect(result.headers['X-RateLimit-Remaining']).toBe(19);
      expect(result.headers['X-RateLimit-Reset']).toBeDefined();
    });

    it('should include Retry-After when denied', () => {
      // Exhaust tokens
      for (let i = 0; i < 20; i++) {
        limiter.check('client1');
      }

      const result = limiter.check('client1');
      expect(result.headers['Retry-After']).toBeDefined();
      expect(result.headers['Retry-After']).toBeGreaterThan(0);
    });
  });
});

describe('SlidingWindowRateLimiter', () => {
  let limiter: SlidingWindowRateLimiter;

  beforeEach(() => {
    limiter = createSlidingWindowLimiter(1000, 10); // 10 requests per second
  });

  afterEach(() => {
    limiter.dispose();
  });

  describe('basic functionality', () => {
    it('should allow requests within limit', () => {
      const result = limiter.check('client1');

      expect(result.allowed).toBe(true);
      expect(result.remaining).toBe(9);
    });

    it('should deny requests exceeding limit', () => {
      // Make 10 requests
      for (let i = 0; i < 10; i++) {
        limiter.check('client1');
      }

      const result = limiter.check('client1');
      expect(result.allowed).toBe(false);
      expect(result.remaining).toBe(0);
    });

    it('should track per-client limits', () => {
      // Use client1's quota
      for (let i = 0; i < 10; i++) {
        limiter.check('client1');
      }

      // client2 should still have quota
      const result = limiter.check('client2');
      expect(result.allowed).toBe(true);
    });
  });

  describe('sliding window behavior', () => {
    it('should allow requests after window slides', async () => {
      vi.useFakeTimers();

      // Use all quota
      for (let i = 0; i < 10; i++) {
        limiter.check('client1');
      }

      // Wait for window to slide
      vi.advanceTimersByTime(1100);

      const result = limiter.check('client1');
      expect(result.allowed).toBe(true);

      vi.useRealTimers();
    });
  });

  describe('reset', () => {
    it('should reset specific client', () => {
      for (let i = 0; i < 5; i++) {
        limiter.check('client1');
      }

      limiter.reset('client1');

      const result = limiter.check('client1');
      expect(result.remaining).toBe(9);
    });

    it('should reset all clients', () => {
      limiter.check('client1');
      limiter.check('client2');

      limiter.reset();

      const result1 = limiter.check('client1');
      const result2 = limiter.check('client2');

      expect(result1.remaining).toBe(9);
      expect(result2.remaining).toBe(9);
    });
  });
});

describe('createStrictRateLimiter', () => {
  it('should create limiter with strict defaults', () => {
    const limiter = createStrictRateLimiter();

    // Should have lower limits (10 req/s, 20 burst)
    const stats = limiter.getStats();
    expect(stats.activeBuckets).toBe(1);

    // Exhaust 20 tokens
    for (let i = 0; i < 20; i++) {
      limiter.check('client1');
    }

    const result = limiter.check('client1');
    expect(result.allowed).toBe(false);

    limiter.dispose();
  });
});