import { rm } from 'fs/promises'; import { AbortPromptError, ExitPromptError } from '@inquirer/core'; import * as yaml from 'js-yaml'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import logger from '../src/logger'; import { createDummyFiles, initializeProject, reportProviderAPIKeyWarnings, } from '../src/onboarding'; import { TestSuiteConfigSchema } from '../src/types/index'; import { mockProcessEnv } from './util/utils'; // Create hoisted mocks for inquirer modules const mockSelect = vi.hoisted(() => vi.fn()); const mockCheckbox = vi.hoisted(() => vi.fn()); const mockConfirm = vi.hoisted(() => vi.fn()); const mockFs = vi.hoisted(() => ({ existsSync: vi.fn(), writeFileSync: vi.fn(), mkdirSync: vi.fn(), access: vi.fn((path: string) => { if (mockFs.existsSync(path)) { return undefined; } throw Object.assign(new Error(`ENOENT: no such file or directory, access '${path}'`), { code: 'ENOENT', }); }), })); vi.mock('fs', () => ({ default: mockFs, ...mockFs, })); vi.mock('fs/promises', () => ({ default: { access: mockFs.access, writeFile: mockFs.writeFileSync, mkdir: mockFs.mkdirSync, mkdtemp: vi.fn(), rm: vi.fn(), }, access: mockFs.access, writeFile: mockFs.writeFileSync, mkdir: mockFs.mkdirSync, mkdtemp: vi.fn(), rm: vi.fn(), })); vi.mock('glob', () => ({ globSync: vi.fn(), })); vi.mock('libsql'); vi.mock('@inquirer/select', () => ({ __esModule: true, default: mockSelect, })); vi.mock('@inquirer/checkbox', () => ({ __esModule: true, default: mockCheckbox, })); vi.mock('@inquirer/confirm', () => ({ __esModule: true, default: mockConfirm, })); vi.mock('../src/database', () => ({ getDb: vi.fn(), })); vi.mock('../src/telemetry', () => ({ default: { record: vi.fn() }, record: vi.fn(), })); vi.mock('../src/util/fetch/index.ts', () => ({ fetch: vi.fn(), })); vi.mock('../src/redteam/commands/init', () => ({ redteamInit: vi.fn(), })); vi.mock('../src/envars', () => ({ getEnvString: vi.fn(), getEnvBool: vi.fn(() => false), getEnvInt: vi.fn((_key: string, defaultValue: number) => defaultValue), })); beforeEach(() => { vi.clearAllMocks(); mockSelect.mockReset(); mockCheckbox.mockReset(); mockConfirm.mockReset(); mockFs.existsSync.mockReset(); mockFs.writeFileSync.mockReset(); mockFs.mkdirSync.mockReset(); mockFs.access.mockClear(); }); describe('reportProviderAPIKeyWarnings', () => { const openaiID = 'openai:gpt-4o'; const anthropicID = 'anthropic:messages:claude-3-5-sonnet-20241022'; let restoreEnv: () => void; beforeEach(() => { restoreEnv = mockProcessEnv({ ANTHROPIC_API_KEY: '', OPENAI_API_KEY: '', }); }); afterEach(() => { restoreEnv(); }); it('should produce a warning for openai if env key is not set', () => { expect(reportProviderAPIKeyWarnings([openaiID])).toEqual( expect.arrayContaining([ expect.stringContaining('OPENAI_API_KEY environment variable is not set'), ]), ); }); it('should produce a warning for anthropic if env key is not set', () => { expect(reportProviderAPIKeyWarnings([anthropicID])).toEqual( expect.arrayContaining([ expect.stringContaining('ANTHROPIC_API_KEY environment variable is not set'), ]), ); }); it('should produce multiple warnings for applicable providers if env keys are not set', () => { expect(reportProviderAPIKeyWarnings([openaiID, anthropicID])).toEqual( expect.arrayContaining([ expect.stringContaining('OPENAI_API_KEY environment variable is not set'), expect.stringContaining('ANTHROPIC_API_KEY environment variable is not set'), ]), ); }); it('should be able to accept an object input so long as it has a valid id field', () => { expect(reportProviderAPIKeyWarnings([{ id: openaiID }, anthropicID])).toEqual( expect.arrayContaining([ expect.stringContaining('OPENAI_API_KEY environment variable is not set'), expect.stringContaining('ANTHROPIC_API_KEY environment variable is not set'), ]), ); }); it('should produce only warnings for applicable providers if the env keys are not set', () => { mockProcessEnv({ OPENAI_API_KEY: '' }); expect(reportProviderAPIKeyWarnings([openaiID, anthropicID])).toEqual( expect.arrayContaining([ expect.stringContaining('ANTHROPIC_API_KEY environment variable is not set'), ]), ); }); }); describe('createDummyFiles', () => { let tempDir: string; beforeEach(() => { tempDir = '/fake/temp/dir'; mockConfirm.mockResolvedValue(true); mockFs.existsSync.mockReturnValue(false); mockFs.writeFileSync.mockImplementation(() => undefined); }); afterEach(async () => { await rm(tempDir, { recursive: true, force: true }); }); it('should generate a valid YAML configuration file that matches TestSuiteConfigSchema', async () => { await createDummyFiles(tempDir, false); const configCall = mockFs.writeFileSync.mock.calls.find((call: any[]) => call[0].toString().endsWith('promptfooconfig.yaml'), ); const readmeCall = mockFs.writeFileSync.mock.calls.find((call: any[]) => call[0].toString().endsWith('README.md'), ); expect(configCall).toBeDefined(); expect(readmeCall).toBeDefined(); const configContent = configCall?.[1] as string; expect(configContent).toBeDefined(); const parsedConfig = yaml.load(configContent); const validationResult = TestSuiteConfigSchema.safeParse(parsedConfig); expect(validationResult.success).toBe(true); // Assert that validation was successful and config is defined expect(validationResult.data).toBeDefined(); const config = validationResult.data!; expect(config.prompts).toHaveLength(2); expect(config.providers).toHaveLength(2); expect(config.providers).toContain('openai:gpt-5-mini'); expect(config.providers).toContain('openai:gpt-5'); }); it('should generate valid YAML configuration for RAG setup', async () => { mockSelect .mockResolvedValueOnce('rag') .mockResolvedValueOnce('python') .mockResolvedValueOnce('openai:gpt-4o'); await createDummyFiles(tempDir, true); const configCall = mockFs.writeFileSync.mock.calls.find((call: any[]) => call[0].toString().endsWith('promptfooconfig.yaml'), ); const contextCall = mockFs.writeFileSync.mock.calls.find((call: any[]) => call[0].toString().endsWith('context.py'), ); expect(configCall).toBeDefined(); expect(contextCall).toBeDefined(); const configContent = configCall?.[1] as string; expect(configContent).toBeDefined(); const parsedConfig = yaml.load(configContent); const validationResult = TestSuiteConfigSchema.safeParse(parsedConfig); expect(validationResult.success).toBe(true); // Assert that validation was successful and config is defined expect(validationResult.data).toBeDefined(); const config = validationResult.data!; expect(config.tests).toBeDefined(); expect(Array.isArray(config.tests)).toBe(true); const tests = config.tests as any[]; expect(tests.length).toBeGreaterThan(0); const firstTest = tests[0]; expect(firstTest).toBeTruthy(); expect(typeof firstTest).toBe('object'); expect(firstTest).toHaveProperty('vars'); const { vars } = firstTest; expect(typeof vars).toBe('object'); expect(vars).toHaveProperty('inquiry'); expect(vars).toHaveProperty('context'); expect(mockSelect).toHaveBeenCalledTimes(3); expect(mockCheckbox).toHaveBeenCalledTimes(0); expect(mockConfirm).toHaveBeenCalledTimes(0); }); it('should prompt for confirmation when files exist', async () => { mockFs.existsSync.mockImplementation((path: string) => path.includes('promptfooconfig.yaml')); mockConfirm.mockResolvedValueOnce(true); mockSelect.mockResolvedValueOnce('compare').mockResolvedValueOnce('openai:gpt-4o'); await createDummyFiles(tempDir, true); expect(mockConfirm).toHaveBeenCalledTimes(1); expect(mockConfirm).toHaveBeenCalledWith( expect.objectContaining({ message: expect.stringContaining('already exist'), }), ); }); }); describe('initializeProject', () => { let originalExitCode: number | string | null | undefined; beforeEach(() => { originalExitCode = process.exitCode; process.exitCode = undefined; }); afterEach(() => { vi.restoreAllMocks(); process.exitCode = originalExitCode; }); it('should return after prompt cancellation without terminating the host process', async () => { mockSelect.mockRejectedValueOnce(new ExitPromptError()); await expect(initializeProject(null, true)).resolves.toBeUndefined(); expect(process.exitCode).toBe(130); }); it('should return after prompt abort without terminating the host process', async () => { mockSelect.mockRejectedValueOnce(new AbortPromptError()); await expect(initializeProject(null, true)).resolves.toBeUndefined(); expect(process.exitCode).toBe(130); }); it('should print current-directory next steps for the redteam path', async () => { // `promptfoo init` with no directory arg, choosing "Run a red team evaluation". mockSelect.mockResolvedValueOnce('redteam'); const infoSpy = vi.spyOn(logger, 'info').mockImplementation(() => logger); await initializeProject(null, true); const messages = infoSpy.mock.calls.map((call) => String(call[0])); // Regression: the redteam branch dropped `outDirectory`, so the "Next steps" // output interpolated `undefined` instead of using the current-directory flow. expect(messages).toEqual( expect.arrayContaining([ expect.stringContaining('Setup complete! Next steps:'), expect.stringContaining('to evaluate your prompts'), expect.stringContaining('to view results in your browser'), ]), ); expect(messages.some((msg) => msg.includes('undefined'))).toBe(false); }); });