import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; import { fetchWithCache } from '../../../src/cache'; import { VERSION } from '../../../src/constants'; import { DEFAULT_PURPOSE, extractSystemPurpose } from '../../../src/redteam/extraction/purpose'; import { getRemoteGenerationUrl } from '../../../src/redteam/remoteGeneration'; import { createMockProvider, createProviderResponse, type MockApiProvider, } from '../../factories/provider'; import { mockProcessEnv } from '../../util/utils'; vi.mock('../../../src/logger', () => ({ default: { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn(), }, })); vi.mock('../../../src/cache', async (importOriginal) => { return { ...(await importOriginal()), fetchWithCache: vi.fn(), }; }); vi.mock('../../../src/redteam/remoteGeneration', async () => ({ ...(await vi.importActual('../../../src/redteam/remoteGeneration')), getRemoteGenerationUrl: vi.fn().mockReturnValue('https://api.promptfoo.app/api/v1/task'), })); describe('System Purpose Extractor', () => { let provider: MockApiProvider; let originalEnv: NodeJS.ProcessEnv; beforeAll(() => { originalEnv = { ...process.env }; }); beforeEach(() => { mockProcessEnv({ ...originalEnv }, { clear: true }); mockProcessEnv({ PROMPTFOO_REMOTE_GENERATION_URL: undefined }); provider = createMockProvider({ response: createProviderResponse({ output: 'Extracted system purpose', }), }); vi.clearAllMocks(); vi.mocked(getRemoteGenerationUrl).mockImplementation(function () { return 'https://api.promptfoo.app/api/v1/task'; }); }); afterEach(() => { mockProcessEnv(originalEnv, { clear: true }); }); it('should use remote generation when enabled', async () => { mockProcessEnv({ OPENAI_API_KEY: undefined }); mockProcessEnv({ PROMPTFOO_DISABLE_REDTEAM_REMOTE_GENERATION: 'false' }); vi.mocked(fetchWithCache).mockResolvedValue({ data: { task: 'purpose', result: 'Remote extracted purpose' }, status: 200, statusText: 'OK', cached: false, }); const result = await extractSystemPurpose(provider, ['prompt1', 'prompt2'], { providerTargetIds: ['file://local-provider.ts'], cloudTargetId: 'cloud-target-123', }); expect(result).toBe('Remote extracted purpose'); expect(fetchWithCache).toHaveBeenCalledWith( 'https://api.promptfoo.app/api/v1/task', expect.objectContaining({ method: 'POST', body: JSON.stringify({ task: 'purpose', prompts: ['prompt1', 'prompt2'], version: VERSION, email: null, targetId: 'cloud-target-123', }), }), expect.any(Number), 'json', ); }); it('should not fall back to local extraction when remote generation fails', async () => { mockProcessEnv({ PROMPTFOO_DISABLE_REDTEAM_REMOTE_GENERATION: 'false' }); const originalOpenaiKey = process.env.OPENAI_API_KEY; mockProcessEnv({ OPENAI_API_KEY: undefined }); vi.mocked(fetchWithCache).mockRejectedValue(new Error('Remote generation failed')); const result = await extractSystemPurpose(provider, ['prompt1', 'prompt2']); expect(result).toBe(''); expect(provider.callApi).not.toHaveBeenCalled(); mockProcessEnv({ OPENAI_API_KEY: originalOpenaiKey }); }); it('should use local extraction when remote generation is disabled', async () => { mockProcessEnv({ PROMPTFOO_DISABLE_REDTEAM_REMOTE_GENERATION: 'true' }); const result = await extractSystemPurpose(provider, ['prompt']); expect(result).toBe('Extracted system purpose'); expect(provider.callApi).toHaveBeenCalledWith(expect.stringContaining('prompt')); expect(fetchWithCache).not.toHaveBeenCalled(); }); it('should extract system purpose when returned without xml tags', async () => { mockProcessEnv({ PROMPTFOO_DISABLE_REDTEAM_REMOTE_GENERATION: 'true' }); vi.mocked(provider.callApi).mockResolvedValue({ output: 'Extracted system purpose' }); const result = await extractSystemPurpose(provider, ['prompt1', 'prompt2']); expect(result).toBe('Extracted system purpose'); }); it('should return default message for empty prompts array', async () => { const result = await extractSystemPurpose(provider, []); expect(result).toBe(DEFAULT_PURPOSE); }); it('should return default message for prompts array with only template variable', async () => { const result = await extractSystemPurpose(provider, ['{{prompt}}']); expect(result).toBe(DEFAULT_PURPOSE); }); });