sentryagent-idp/tests/unit/middleware/rateLimit.test.ts

/**
 * Unit tests for src/middleware/rateLimit.ts
 */

import { Request, Response, NextFunction } from 'express';
import { RateLimitError } from '../../../src/utils/errors';

const mockIncr = jest.fn();
const mockExpire = jest.fn();

jest.mock('../../../src/cache/redis', () => ({
  getRedisClient: jest.fn().mockResolvedValue({
    incr: mockIncr,
    expire: mockExpire,
  }),
}));

import { rateLimitMiddleware } from '../../../src/middleware/rateLimit';

function buildMocks(clientId?: string): {
  req: Partial<Request>;
  res: Partial<Response>;
  next: NextFunction;
} {
  const res: Partial<Response> = {
    setHeader: jest.fn(),
  };
  return {
    req: {
      user: clientId ? { client_id: clientId, sub: clientId, scope: '', jti: '', iat: 0, exp: 0 } : undefined,
      ip: '127.0.0.1',
    },
    res,
    next: jest.fn() as NextFunction,
  };
}

describe('rateLimitMiddleware', () => {
  beforeEach(() => {
    jest.clearAllMocks();
    mockExpire.mockResolvedValue(1);
  });

  it('should set X-RateLimit-* headers and call next() when counter is under the limit', async () => {
    mockIncr.mockResolvedValue(1);
    const { req, res, next } = buildMocks('agent-123');

    await rateLimitMiddleware(req as Request, res as Response, next);

    expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Limit', 100);
    expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Remaining', 99);
    expect(res.setHeader).toHaveBeenCalledWith('X-RateLimit-Reset', expect.any(Number));
    expect(next).toHaveBeenCalledWith();
    expect(next).not.toHaveBeenCalledWith(expect.any(Error));
  });

  it('should call next(RateLimitError) when counter equals 100', async () => {
    mockIncr.mockResolvedValue(101);
    const { req, res, next } = buildMocks('agent-456');

    await rateLimitMiddleware(req as Request, res as Response, next);

    expect(next).toHaveBeenCalledWith(expect.any(RateLimitError));
  });

  it('should use req.ip as key when req.user is not set', async () => {
    mockIncr.mockResolvedValue(5);
    const { req, res, next } = buildMocks(); // no clientId → no req.user

    await rateLimitMiddleware(req as Request, res as Response, next);

    expect(mockIncr).toHaveBeenCalledWith(expect.stringContaining('127.0.0.1'));
    expect(next).toHaveBeenCalledWith();
  });

  it('should set expire TTL only on first request (count === 1)', async () => {
    mockIncr.mockResolvedValue(1);
    const { req, res, next } = buildMocks('agent-789');

    await rateLimitMiddleware(req as Request, res as Response, next);

    expect(mockExpire).toHaveBeenCalledWith(expect.any(String), 60);
  });

  it('should not call expire on subsequent requests (count > 1)', async () => {
    mockIncr.mockResolvedValue(50);
    const { req, res, next } = buildMocks('agent-789');

    await rateLimitMiddleware(req as Request, res as Response, next);

    expect(mockExpire).not.toHaveBeenCalled();
  });
});