promptfoo--promptfoo/test/main.test.ts

import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { pathToFileURL } from 'node:url';

import { Command } from 'commander';
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';

// Hoisted mocks for shutdown tests
const mockSetupEnv = vi.hoisted(() => vi.fn());
const mockSetLogLevel = vi.hoisted(() => vi.fn());
const mockTelemetryRecord = vi.hoisted(() => vi.fn());
const mockTelemetryInitialize = vi.hoisted(() => vi.fn());
const mockTelemetryShutdown = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
const mockCloseLogger = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
const mockCloseDbIfOpen = vi.hoisted(() => vi.fn());
const mockDispatcherDestroy = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
const mockGetGlobalDispatcher = vi.hoisted(() =>
  vi.fn().mockReturnValue({ destroy: mockDispatcherDestroy }),
);
// Import actual undici to preserve other exports (Agent, ProxyAgent, setGlobalDispatcher, etc.)
const actualUndici = vi.hoisted(() => vi.importActual<typeof import('undici')>('undici'));

// Mock the dependencies
vi.mock('../src/util', () => ({
  setupEnv: mockSetupEnv,
}));

vi.mock('../src/logger', () => ({
  __esModule: true,
  default: { debug: vi.fn(), warn: vi.fn(), error: vi.fn(), info: vi.fn() },
  setLogLevel: mockSetLogLevel,
  closeLogger: mockCloseLogger,
}));

vi.mock('../src/telemetry', () => ({
  __esModule: true,
  default: {
    initialize: mockTelemetryInitialize,
    record: mockTelemetryRecord,
    shutdown: mockTelemetryShutdown,
  },
}));

vi.mock('../src/database/index', () => ({
  closeDbIfOpen: mockCloseDbIfOpen,
}));

vi.mock('undici', async () => ({
  ...(await actualUndici),
  getGlobalDispatcher: mockGetGlobalDispatcher,
}));

// Mock code scan commands to avoid ESM import issues with execa
vi.mock('../src/codeScan', () => ({
  codeScansCommand: vi.fn(),
}));

let addCommonOptionsRecursively: typeof import('../src/mainUtils').addCommonOptionsRecursively;
let isMainModule: typeof import('../src/mainUtils').isMainModule;
let shouldSkipDefaultConfigLoading: typeof import('../src/mainUtils').shouldSkipDefaultConfigLoading;
let setupEnvFilesFromArgv: typeof import('../src/mainUtils').setupEnvFilesFromArgv;
let shutdownGracefully: typeof import('../src/mainUtils').shutdownGracefully;

async function loadMainModule() {
  vi.resetModules();
  ({
    addCommonOptionsRecursively,
    isMainModule,
    shouldSkipDefaultConfigLoading,
    setupEnvFilesFromArgv,
    shutdownGracefully,
  } = await import('../src/mainUtils'));
}

describe('setupEnvFilesFromArgv', () => {
  beforeEach(async () => {
    await loadMainModule();
    mockSetupEnv.mockReset();
    mockTelemetryInitialize.mockReset();
  });

  it('should load env files before command actions run', () => {
    setupEnvFilesFromArgv(['eval', '--env-file', '.env.local']);

    expect(mockSetupEnv).toHaveBeenCalledWith('.env.local', { refreshConfigDirectory: true });
    expect(mockTelemetryInitialize).toHaveBeenCalledOnce();
  });

  it('should support repeated and comma-separated env file args', () => {
    setupEnvFilesFromArgv(['eval', '--env-file', '.env.one', '--env-path=.env.two,.env.three']);

    expect(mockSetupEnv).toHaveBeenCalledWith(['.env.one', '.env.two', '.env.three'], {
      refreshConfigDirectory: true,
    });
  });

  it('should ignore flags after --', () => {
    setupEnvFilesFromArgv(['eval', '--', '--env-file', '.env.local']);

    expect(mockSetupEnv).not.toHaveBeenCalled();
    expect(mockTelemetryInitialize).toHaveBeenCalledOnce();
  });

  it('should recognize the --env-path alias', () => {
    setupEnvFilesFromArgv(['eval', '--env-path', '.env.staging']);

    expect(mockSetupEnv).toHaveBeenCalledWith('.env.staging', { refreshConfigDirectory: true });
  });

  it('should be a no-op when no env flags are present', () => {
    setupEnvFilesFromArgv(['eval', '--verbose', '--no-cache']);

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

  it('should ignore --env-file= with empty value', () => {
    setupEnvFilesFromArgv(['eval', '--env-file=']);

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

  it('should skip --env-file when next arg starts with -', () => {
    setupEnvFilesFromArgv(['eval', '--env-file', '--verbose']);

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

describe('shouldSkipDefaultConfigLoading', () => {
  beforeEach(async () => {
    await loadMainModule();
  });

  it('skips unrelated default config discovery for code-scans commands', () => {
    expect(shouldSkipDefaultConfigLoading(['code-scans', 'run'])).toBe(true);
    expect(shouldSkipDefaultConfigLoading(['--verbose', 'code-scans', 'run', '--help'])).toBe(true);
    expect(
      shouldSkipDefaultConfigLoading(['--env-file', '.env.local', 'code-scans', 'run', '--help']),
    ).toBe(true);
    expect(shouldSkipDefaultConfigLoading(['--env-path=.env.local', 'code-scans', 'run'])).toBe(
      true,
    );
  });

  it('keeps default config discovery for other commands and post-separator arguments', () => {
    expect(shouldSkipDefaultConfigLoading(['eval', '--help'])).toBe(false);
    expect(shouldSkipDefaultConfigLoading(['--', 'code-scans', 'run'])).toBe(false);
  });
});

describe('addCommonOptionsRecursively', () => {
  const originalExit = process.exit;
  let program: Command;
  let subCommand: Command;

  beforeAll(() => {
    process.exit = vi.fn() as any;
  });

  beforeEach(async () => {
    await loadMainModule();
    program = new Command('promptfoo');
    program.action(() => {});
    subCommand = program.command('subcommand');
    subCommand.action(() => {});
    mockTelemetryRecord.mockReset();
    vi.clearAllMocks();
  });

  afterAll(() => {
    process.exit = originalExit;
  });

  it('should add both verbose and env-file options to a command', () => {
    addCommonOptionsRecursively(program);

    const hasVerboseOption = program.options.some(
      (option) => option.short === '-v' || option.long === '--verbose',
    );
    const hasEnvFileOption = program.options.some(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    );

    expect(hasVerboseOption).toBe(true);
    expect(hasEnvFileOption).toBe(true);
  });

  it('should not add duplicate options if they already exist', () => {
    program.option('--env-file, --env-path <path>', 'Path to .env file');
    program.option('-v, --verbose', 'Show debug logs', false);

    // Count options before
    const envFileOptionsBefore = program.options.filter(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    ).length;

    const verboseOptionsBefore = program.options.filter(
      (option) => option.short === '-v' || option.long === '--verbose',
    ).length;

    addCommonOptionsRecursively(program);

    // Count options after
    const envFileOptionsAfter = program.options.filter(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    ).length;

    const verboseOptionsAfter = program.options.filter(
      (option) => option.short === '-v' || option.long === '--verbose',
    ).length;

    // Should still have the same number of options
    expect(envFileOptionsAfter).toBe(envFileOptionsBefore);
    expect(verboseOptionsAfter).toBe(verboseOptionsBefore);
  });

  it('should add options to subcommands', () => {
    addCommonOptionsRecursively(program);

    const hasSubcommandVerboseOption = subCommand.options.some(
      (option) => option.short === '-v' || option.long === '--verbose',
    );
    const hasSubcommandEnvFileOption = subCommand.options.some(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    );

    expect(hasSubcommandVerboseOption).toBe(true);
    expect(hasSubcommandEnvFileOption).toBe(true);
  });

  it('should add options to nested subcommands', () => {
    // Create a deeper command structure
    const subSubCommand = subCommand.command('subsubcommand');
    subSubCommand.action(() => {});
    const level3Command = subSubCommand.command('subsubsubcommand');
    level3Command.action(() => {});

    addCommonOptionsRecursively(program);

    // Check all levels of commands have the options
    const hasMainVerboseOption = program.options.some(
      (option) => option.short === '-v' || option.long === '--verbose',
    );
    const hasMainEnvFileOption = program.options.some(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    );

    const hasSubCommandVerboseOption = subCommand.options.some(
      (option) => option.short === '-v' || option.long === '--verbose',
    );
    const hasSubCommandEnvFileOption = subCommand.options.some(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    );

    const hasSubSubCommandVerboseOption = subSubCommand.options.some(
      (option) => option.short === '-v' || option.long === '--verbose',
    );
    const hasSubSubCommandEnvFileOption = subSubCommand.options.some(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    );

    const hasSubSubSubCommandVerboseOption = level3Command.options.some(
      (option) => option.short === '-v' || option.long === '--verbose',
    );
    const hasSubSubSubCommandEnvFileOption = level3Command.options.some(
      (option) => option.long === '--env-file' || option.long === '--env-path',
    );

    expect(hasMainVerboseOption).toBe(true);
    expect(hasMainEnvFileOption).toBe(true);
    expect(hasSubCommandVerboseOption).toBe(true);
    expect(hasSubCommandEnvFileOption).toBe(true);
    expect(hasSubSubCommandVerboseOption).toBe(true);
    expect(hasSubSubCommandEnvFileOption).toBe(true);
    expect(hasSubSubSubCommandVerboseOption).toBe(true);
    expect(hasSubSubSubCommandEnvFileOption).toBe(true);
  });

  it('should register a single hook that handles both options', () => {
    // Create a fake action that manually mocks the Commander hook system
    const mockHookRegister = vi.fn();
    (program as any).hook = mockHookRegister;

    // Apply common options
    addCommonOptionsRecursively(program);

    // Verify the hook was registered only once
    expect(mockHookRegister).toHaveBeenCalledTimes(1);
    expect(mockHookRegister).toHaveBeenCalledWith('preAction', expect.any(Function));

    // Get the hook function
    const preActionFn = mockHookRegister.mock.calls[0][1];

    // Create mock command object with name() and parent for telemetry
    const createMockCommand = (opts: Record<string, unknown>, commandName = 'test') => ({
      opts: () => opts,
      name: () => commandName,
      parent: null,
    });

    // Test verbose option
    preActionFn(createMockCommand({ verbose: true }));
    expect(mockSetLogLevel).toHaveBeenCalledWith('debug');

    // Test env-file option
    preActionFn(createMockCommand({ envFile: '.env.test' }));
    expect(mockSetupEnv).toHaveBeenCalledWith('.env.test');

    // Test both options together
    preActionFn(createMockCommand({ verbose: true, envFile: '.env.combined' }));
    expect(mockSetLogLevel).toHaveBeenCalledWith('debug');
    expect(mockSetupEnv).toHaveBeenCalledWith('.env.combined');
  });

  it('keeps structured code-scan output muted even when --verbose is present', () => {
    const mockHookRegister = vi.fn();
    (program as any).hook = mockHookRegister;

    addCommonOptionsRecursively(program);

    const preActionFn = mockHookRegister.mock.calls[0][1];
    const originalArgv = process.argv;

    process.argv = ['node', 'promptfoo', 'code-scans', 'run', '.', '--json', '--verbose'];

    try {
      preActionFn({
        opts: () => ({ verbose: true }),
        name: () => 'run',
        parent: {
          name: () => 'code-scans',
          parent: null,
        },
      });
    } finally {
      process.argv = originalArgv;
    }

    expect(mockSetLogLevel).not.toHaveBeenCalledWith('debug');
  });

  it('should parse --env-file without consuming positional subcommand arguments', async () => {
    const action = vi.fn();
    const documentCommand = program.command('scan-model').argument('<model>').action(action);

    addCommonOptionsRecursively(program);

    await program.parseAsync(['scan-model', '--env-file', '.env.local', 'llama-3'], {
      from: 'user',
    });

    expect(action).toHaveBeenCalledWith(
      'llama-3',
      expect.objectContaining({ verbose: false }),
      documentCommand,
    );
    expect(mockSetupEnv).toHaveBeenCalledWith('.env.local');
  });

  it('should support repeated and comma-separated --env-file values', async () => {
    addCommonOptionsRecursively(program);

    await program.parseAsync(
      ['subcommand', '--env-file', '.env.one', '--env-file', '.env.two,.env.three'],
      { from: 'user' },
    );

    expect(mockSetupEnv).toHaveBeenCalledWith(['.env.one', '.env.two', '.env.three']);
  });

  it('should record telemetry once for nested subcommands', async () => {
    const nestedParent = program.command('redteam');
    const nestedAction = vi.fn();
    nestedParent.command('run').action(nestedAction);

    addCommonOptionsRecursively(program);

    await program.parseAsync(['redteam', 'run'], { from: 'user' });

    expect(nestedAction).toHaveBeenCalledTimes(1);
    expect(mockTelemetryRecord).toHaveBeenCalledTimes(1);
    expect(mockTelemetryRecord).toHaveBeenCalledWith('command_used', { name: 'redteam run' });
  });
});

describe('isMainModule', () => {
  let tempDir: string;
  let realFilePath: string;
  let symlinkPath: string;

  beforeAll(async () => {
    await loadMainModule();
    // Create a temporary directory with a real file and a symlink
    tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'main-test-'));
    realFilePath = path.join(tempDir, 'real-file.js');
    symlinkPath = path.join(tempDir, 'symlink-file.js');

    // Create a real file
    fs.writeFileSync(realFilePath, '// test file');

    // Create a symlink pointing to the real file
    fs.symlinkSync(realFilePath, symlinkPath);
  });

  afterAll(() => {
    // Clean up temporary files
    try {
      fs.unlinkSync(symlinkPath);
      fs.unlinkSync(realFilePath);
      fs.rmdirSync(tempDir);
    } catch {
      // Ignore cleanup errors
    }
  });

  it('should return false when processArgv1 is undefined', () => {
    const result = isMainModule('file:///some/path/main.js', undefined);
    expect(result).toBe(false);
  });

  it('should return false when processArgv1 is empty string', () => {
    const result = isMainModule('file:///some/path/main.js', '');
    expect(result).toBe(false);
  });

  it('should return true when paths match directly', () => {
    const fileUrl = pathToFileURL(realFilePath).href;
    const result = isMainModule(fileUrl, realFilePath);
    expect(result).toBe(true);
  });

  it('should return true when processArgv1 is a symlink pointing to the module', () => {
    // This is the key test case for npm global bin symlinks
    const fileUrl = pathToFileURL(realFilePath).href;
    const result = isMainModule(fileUrl, symlinkPath);
    expect(result).toBe(true);
  });

  it('should return false when paths do not match', () => {
    const fileUrl = pathToFileURL(realFilePath).href;
    const otherPath = path.join(tempDir, 'other-file.js');
    fs.writeFileSync(otherPath, '// other file');

    try {
      const result = isMainModule(fileUrl, otherPath);
      expect(result).toBe(false);
    } finally {
      fs.unlinkSync(otherPath);
    }
  });

  it('should return false when processArgv1 points to non-existent file', () => {
    const fileUrl = pathToFileURL(realFilePath).href;
    const nonExistentPath = path.join(tempDir, 'non-existent.js');
    const result = isMainModule(fileUrl, nonExistentPath);
    expect(result).toBe(false);
  });

  it('should handle relative paths correctly', () => {
    // Save current directory
    const originalCwd = process.cwd();

    try {
      // Change to temp directory
      process.chdir(tempDir);

      const fileUrl = pathToFileURL(realFilePath).href;
      const result = isMainModule(fileUrl, './real-file.js');
      expect(result).toBe(true);
    } finally {
      // Restore original directory
      process.chdir(originalCwd);
    }
  });

  it('should handle symlink to symlink chains', () => {
    const secondSymlinkPath = path.join(tempDir, 'second-symlink.js');

    try {
      // Create a symlink pointing to the first symlink
      fs.symlinkSync(symlinkPath, secondSymlinkPath);

      const fileUrl = pathToFileURL(realFilePath).href;
      const result = isMainModule(fileUrl, secondSymlinkPath);
      expect(result).toBe(true);
    } finally {
      try {
        fs.unlinkSync(secondSymlinkPath);
      } catch {
        // Ignore cleanup errors
      }
    }
  });
});

describe('shutdownGracefully', () => {
  const originalExit = process.exit;

  beforeEach(async () => {
    await loadMainModule();
    vi.useFakeTimers();
    process.exit = vi.fn() as never;
    // Reset all mocks to default behavior
    mockSetupEnv.mockReset();
    mockSetLogLevel.mockReset();
    mockTelemetryShutdown.mockReset().mockResolvedValue(undefined);
    mockCloseLogger.mockReset().mockResolvedValue(undefined);
    mockCloseDbIfOpen.mockReset();
    mockDispatcherDestroy.mockReset().mockResolvedValue(undefined);
    mockGetGlobalDispatcher.mockReset().mockReturnValue({ destroy: mockDispatcherDestroy });
  });

  afterEach(() => {
    vi.useRealTimers();
    vi.resetAllMocks();
    process.exit = originalExit;
  });

  it('should complete gracefully when all operations succeed quickly', async () => {
    const shutdownPromise = shutdownGracefully();

    // Advance past any internal timeouts
    await vi.runAllTimersAsync();

    await shutdownPromise;

    expect(mockTelemetryShutdown).toHaveBeenCalled();
    expect(mockCloseLogger).toHaveBeenCalled();
    expect(mockCloseDbIfOpen).toHaveBeenCalled();
    expect(mockDispatcherDestroy).toHaveBeenCalled();
  });

  it('should handle telemetry shutdown timeout', async () => {
    // Make telemetry.shutdown() hang forever
    mockTelemetryShutdown.mockImplementation(() => new Promise(() => {}));

    const shutdownPromise = shutdownGracefully();

    // Advance time to trigger the individual operation timeout (1s)
    await vi.advanceTimersByTimeAsync(1100);

    // Advance remaining timers
    await vi.runAllTimersAsync();

    await shutdownPromise;

    // Should still have called other cleanup operations
    expect(mockCloseLogger).toHaveBeenCalled();
    expect(mockCloseDbIfOpen).toHaveBeenCalled();
  });

  it('should handle closeLogger timeout', async () => {
    // Make closeLogger() hang forever
    mockCloseLogger.mockImplementation(() => new Promise(() => {}));

    const shutdownPromise = shutdownGracefully();

    // Advance time to trigger all timeouts
    await vi.runAllTimersAsync();

    await shutdownPromise;

    // Other operations should still complete
    expect(mockTelemetryShutdown).toHaveBeenCalled();
    expect(mockCloseDbIfOpen).toHaveBeenCalled();
    expect(mockDispatcherDestroy).toHaveBeenCalled();
  });

  it('should handle dispatcher.destroy() timeout', async () => {
    // Make dispatcher.destroy() hang forever
    mockDispatcherDestroy.mockImplementation(() => new Promise(() => {}));

    const shutdownPromise = shutdownGracefully();

    // Advance time to trigger all timeouts
    await vi.runAllTimersAsync();

    await shutdownPromise;

    // Other operations should still complete
    expect(mockTelemetryShutdown).toHaveBeenCalled();
    expect(mockCloseLogger).toHaveBeenCalled();
    expect(mockCloseDbIfOpen).toHaveBeenCalled();
  });

  it('should handle telemetry shutdown throwing error', async () => {
    mockTelemetryShutdown.mockRejectedValue(new Error('Telemetry failed'));

    const shutdownPromise = shutdownGracefully();

    await vi.runAllTimersAsync();

    // Should not throw - should handle error gracefully
    await expect(shutdownPromise).resolves.toBeUndefined();

    // Other operations should still complete
    expect(mockCloseLogger).toHaveBeenCalled();
    expect(mockCloseDbIfOpen).toHaveBeenCalled();
  });

  it('should handle dispatcher.destroy() throwing error', async () => {
    mockDispatcherDestroy.mockRejectedValue(new Error('Dispatcher failed'));

    const shutdownPromise = shutdownGracefully();

    await vi.runAllTimersAsync();

    // Should not throw
    await expect(shutdownPromise).resolves.toBeUndefined();
  });

  it('should force exit when all cleanup operations hang', async () => {
    // Make all operations hang
    mockTelemetryShutdown.mockImplementation(() => new Promise(() => {}));
    mockCloseLogger.mockImplementation(() => new Promise(() => {}));
    mockDispatcherDestroy.mockImplementation(() => new Promise(() => {}));

    // Start shutdown but don't await yet
    void shutdownGracefully();

    // The force exit timeout is 3000ms
    await vi.advanceTimersByTimeAsync(3000);

    expect(process.exit).toHaveBeenCalledWith(0);
  });

  it('should clear force exit timeout when cleanup completes normally', async () => {
    const setTimeoutSpy = vi.spyOn(global, 'setTimeout');
    const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
    const shutdownPromise = shutdownGracefully();

    await vi.runAllTimersAsync();

    await shutdownPromise;

    const forceExitTimeout = setTimeoutSpy.mock.results[0]?.value;
    expect(clearTimeoutSpy).toHaveBeenCalledWith(forceExitTimeout);
    expect(process.exit).toHaveBeenCalled();
  });

  it('should handle getGlobalDispatcher throwing', async () => {
    mockGetGlobalDispatcher.mockImplementation(() => {
      throw new Error('No dispatcher');
    });

    const shutdownPromise = shutdownGracefully();

    await vi.runAllTimersAsync();

    // Should complete without throwing
    await expect(shutdownPromise).resolves.toBeUndefined();
  });
});