chore: import upstream snapshot with attribution

test/commands/generate/assertions.test.ts

@@ -0,0 +1,233 @@
+import fs from 'fs';
+
+import * as yaml from 'js-yaml';
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import { synthesizeFromTestSuite } from '../../../src/assertions/synthesis';
+import { disableCache } from '../../../src/cache';
+import { doGenerateAssertions } from '../../../src/commands/generate/assertions';
+import telemetry from '../../../src/telemetry';
+import { resolveConfigs } from '../../../src/util/config/load';
+import { loadYaml } from '../../../src/util/yamlLoad';
+
+import type { Assertion, TestSuite } from '../../../src/types/index';
+
+const fsMocks = vi.hoisted(() => ({
+  readFileSync: vi.fn(),
+  writeFileSync: vi.fn(),
+  existsSync: vi.fn(),
+}));
+
+vi.mock('fs', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('fs')>();
+  return {
+    ...actual,
+    default: {
+      ...actual,
+      ...fsMocks,
+    },
+    ...fsMocks,
+  };
+});
+
+vi.mock('fs/promises', () => ({
+  default: {
+    readFile: fsMocks.readFileSync,
+    writeFile: fsMocks.writeFileSync,
+  },
+  readFile: fsMocks.readFileSync,
+  writeFile: fsMocks.writeFileSync,
+}));
+vi.mock('js-yaml');
+vi.mock('../../../src/util/yamlLoad');
+vi.mock('../../../src/assertions/synthesis');
+vi.mock('../../../src/util/config/load', () => ({
+  resolveConfigs: vi.fn(),
+}));
+vi.mock('../../../src/cache', () => ({
+  disableCache: vi.fn(),
+}));
+vi.mock('../../../src/logger', () => ({
+  default: {
+    info: vi.fn(),
+    warn: vi.fn(),
+    error: vi.fn(),
+    debug: vi.fn(),
+    child: vi.fn().mockReturnValue({}),
+  },
+}));
+vi.mock('../../../src/telemetry', () => ({
+  default: {
+    record: vi.fn(),
+    send: vi.fn().mockResolvedValue(undefined),
+  },
+}));
+vi.mock('../../../src/util', () => ({
+  printBorder: vi.fn(),
+  setupEnv: vi.fn(),
+}));
+
+vi.mock('../../../src/util/promptfooCommand', () => ({
+  promptfooCommand: vi.fn().mockReturnValue('promptfoo eval'),
+  detectInstaller: vi.fn().mockReturnValue('unknown'),
+  isRunningUnderNpx: vi.fn().mockReturnValue(false),
+}));
+
+describe('assertion generation', () => {
+  beforeAll(() => {
+    vi.useFakeTimers();
+  });
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+  });
+
+  afterEach(() => {
+    vi.resetAllMocks();
+  });
+
+  afterAll(() => {
+    vi.useRealTimers();
+  });
+
+  describe('doGenerateAssertions', () => {
+    const mockTestSuite: TestSuite = {
+      prompts: [{ raw: 'test prompt', label: 'test' }],
+      tests: [
+        {
+          assert: [
+            {
+              type: 'llm-rubric',
+              value: 'test question',
+            },
+          ],
+        },
+      ],
+      providers: [
+        {
+          id: () => 'test-provider',
+          callApi: vi.fn() as any,
+        },
+      ],
+    };
+
+    const mockResults: Assertion[] = [
+      { type: 'pi', value: 'additional assertion' },
+      { type: 'pi', value: 'additional assertion 2' },
+    ];
+
+    beforeEach(() => {
+      vi.mocked(synthesizeFromTestSuite).mockResolvedValue(mockResults);
+      vi.mocked(yaml.dump).mockReturnValue('yaml content');
+      vi.mocked(loadYaml).mockReturnValue(mockTestSuite);
+      vi.mocked(fs.readFileSync).mockReturnValue('mock config content');
+      vi.mocked(fs.existsSync).mockReturnValue(true);
+      vi.mocked(fs.writeFileSync).mockImplementation(() => undefined);
+      vi.mocked(disableCache).mockImplementation(() => undefined);
+      vi.mocked(telemetry.record).mockImplementation(() => undefined);
+
+      vi.mocked(resolveConfigs).mockResolvedValue({
+        testSuite: mockTestSuite,
+        config: {},
+        basePath: '',
+      });
+    });
+
+    it('should write YAML output', async () => {
+      const configPath = 'config.yaml';
+
+      await doGenerateAssertions({
+        cache: true,
+        config: configPath,
+        numAssertions: '2',
+        type: 'pi',
+        output: 'output.yaml',
+        write: false,
+        defaultConfig: {},
+        defaultConfigPath: configPath,
+      });
+
+      expect(fs.writeFileSync).toHaveBeenCalledWith('output.yaml', 'yaml content');
+    });
+
+    it('should throw error for unsupported file type', async () => {
+      const configPath = 'config.yaml';
+
+      await expect(
+        doGenerateAssertions({
+          cache: true,
+          config: configPath,
+          numAssertions: '2',
+          type: 'pi',
+          output: 'output.txt',
+          write: false,
+          defaultConfig: {},
+          defaultConfigPath: configPath,
+        }),
+      ).rejects.toThrow('Unsupported output file type: output.txt');
+    });
+
+    it('should write to config file when write option is true', async () => {
+      const configPath = 'config.yaml';
+
+      await doGenerateAssertions({
+        cache: true,
+        config: configPath,
+        numAssertions: '2',
+        type: 'pi',
+        write: true,
+        defaultConfig: {},
+        defaultConfigPath: configPath,
+      });
+
+      expect(fs.writeFileSync).toHaveBeenCalledWith(configPath, expect.any(String));
+    });
+
+    it('should throw error when no config file found', async () => {
+      await expect(
+        doGenerateAssertions({
+          cache: true,
+          numAssertions: '2',
+          type: 'pi',
+          write: false,
+          defaultConfig: {},
+          defaultConfigPath: undefined,
+        }),
+      ).rejects.toThrow('Could not find a config file');
+    });
+
+    it('should handle output without file extension', async () => {
+      const configPath = 'config.yaml';
+
+      await expect(
+        doGenerateAssertions({
+          cache: true,
+          config: configPath,
+          numAssertions: '2',
+          type: 'pi',
+          output: 'output',
+          write: false,
+          defaultConfig: {},
+          defaultConfigPath: configPath,
+        }),
+      ).rejects.toThrow('Unsupported output file type: output');
+    });
+
+    it('should handle synthesis errors', async () => {
+      vi.mocked(synthesizeFromTestSuite).mockRejectedValue(new Error('Synthesis failed'));
+      const configPath = 'config.yaml';
+
+      await expect(
+        doGenerateAssertions({
+          cache: true,
+          config: configPath,
+          numAssertions: '2',
+          type: 'pi',
+          output: 'output.yaml',
+          write: false,
+          defaultConfig: {},
+          defaultConfigPath: configPath,
+        }),
+      ).rejects.toThrow('Synthesis failed');
+    });
+  });
+});