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

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

import { Request, Response, NextFunction } from 'express';
import { Pool } from 'pg';
import { ITokenPayload } from '../../../src/types/index';

// Mock pg Pool
jest.mock('pg', () => {
  const mockQuery = jest.fn();
  return {
    Pool: jest.fn().mockImplementation(() => ({ query: mockQuery })),
  };
});

import { createOrgContextMiddleware } from '../../../src/middleware/orgContext';

/** Builds a minimal ITokenPayload for test requests. */
function makeUser(overrides: Partial<ITokenPayload> = {}): ITokenPayload {
  return {
    sub: 'agent-abc-123',
    client_id: 'agent-abc-123',
    scope: 'agents:read',
    jti: 'jti-001',
    iat: 1000,
    exp: 9999999999,
    ...overrides,
  };
}

describe('createOrgContextMiddleware', () => {
  let pool: jest.Mocked<Pool>;
  let mockQuery: jest.Mock;
  let next: jest.MockedFunction<NextFunction>;
  const originalDefaultOrgId = process.env['DEFAULT_ORG_ID'];

  beforeEach(() => {
    jest.clearAllMocks();
    // Get the mocked pool instance and its query function
    pool = new Pool() as jest.Mocked<Pool>;
    mockQuery = pool.query as jest.Mock;
    mockQuery.mockResolvedValue({ rows: [], rowCount: 0 });
    next = jest.fn();
  });

  afterEach(() => {
    if (originalDefaultOrgId === undefined) {
      delete process.env['DEFAULT_ORG_ID'];
    } else {
      process.env['DEFAULT_ORG_ID'] = originalDefaultOrgId;
    }
  });

  it('should set app.organization_id from req.user.organization_id when present', async () => {
    const middleware = createOrgContextMiddleware(pool);
    const req = {
      user: makeUser({ organization_id: 'org_TENANT123' }),
    } as Request;
    const res = {} as Response;

    await middleware(req, res, next);

    expect(mockQuery).toHaveBeenCalledWith('SET app.organization_id = $1', ['org_TENANT123']);
    expect(next).toHaveBeenCalledWith();
    expect(next).toHaveBeenCalledTimes(1);
  });

  it('should fall back to DEFAULT_ORG_ID env var when req.user has no organization_id', async () => {
    process.env['DEFAULT_ORG_ID'] = 'org_default_from_env';
    const middleware = createOrgContextMiddleware(pool);
    const req = {
      user: makeUser({ organization_id: undefined }),
    } as Request;
    const res = {} as Response;

    await middleware(req, res, next);

    expect(mockQuery).toHaveBeenCalledWith('SET app.organization_id = $1', ['org_default_from_env']);
    expect(next).toHaveBeenCalledWith();
  });

  it('should fall back to "org_system" when req.user is absent and DEFAULT_ORG_ID is not set', async () => {
    delete process.env['DEFAULT_ORG_ID'];
    const middleware = createOrgContextMiddleware(pool);
    const req = {} as Request;
    const res = {} as Response;

    await middleware(req, res, next);

    expect(mockQuery).toHaveBeenCalledWith('SET app.organization_id = $1', ['org_system']);
    expect(next).toHaveBeenCalledWith();
  });

  it('should fall back to "org_system" when DEFAULT_ORG_ID is not set and user has no organization_id', async () => {
    delete process.env['DEFAULT_ORG_ID'];
    const middleware = createOrgContextMiddleware(pool);
    const req = {
      user: makeUser({ organization_id: undefined }),
    } as Request;
    const res = {} as Response;

    await middleware(req, res, next);

    expect(mockQuery).toHaveBeenCalledWith('SET app.organization_id = $1', ['org_system']);
    expect(next).toHaveBeenCalledWith();
  });

  it('should call next(err) when pool.query throws', async () => {
    const dbError = new Error('Database connection failed');
    mockQuery.mockRejectedValue(dbError);
    const middleware = createOrgContextMiddleware(pool);
    const req = {
      user: makeUser({ organization_id: 'org_TENANT123' }),
    } as Request;
    const res = {} as Response;

    await middleware(req, res, next);

    expect(next).toHaveBeenCalledWith(dbError);
    expect(next).toHaveBeenCalledTimes(1);
  });

  it('should call next(err) with pool error when req.user is absent and pool throws', async () => {
    const dbError = new Error('Pool exhausted');
    mockQuery.mockRejectedValue(dbError);
    delete process.env['DEFAULT_ORG_ID'];
    const middleware = createOrgContextMiddleware(pool);
    const req = {} as Request;
    const res = {} as Response;

    await middleware(req, res, next);

    expect(next).toHaveBeenCalledWith(dbError);
  });
});