promptfoo--promptfoo/test/util/fileReference.test.ts

import fs from 'fs';
import path from 'path';

import { beforeEach, describe, expect, it, vi } from 'vitest';
import * as esmModule from '../../src/esm';
import logger from '../../src/logger';
import * as pythonUtils from '../../src/python/pythonUtils';
import { isJavascriptFile } from '../../src/util/fileExtensions';
import { loadFileReference, processConfigFileReferences } from '../../src/util/fileReference';
import { loadYaml } from '../../src/util/yamlLoad';
import type { Logger } from 'winston';

vi.mock('fs', async () => {
  const actual = await vi.importActual<typeof import('fs')>('fs');
  return {
    ...actual,
    default: {
      ...actual,
      promises: {
        readFile: vi.fn(),
      },
    },
    promises: {
      readFile: vi.fn(),
    },
  };
});
vi.mock('path');
vi.mock('../../src/util/yamlLoad');
vi.mock('../../src/esm', () => ({
  importModule: vi.fn(),
}));
vi.mock('../../src/python/pythonUtils', () => ({
  runPython: vi.fn(),
}));
vi.mock('../../src/util/fileExtensions');
vi.mock('../../src/logger', () => ({
  default: {
    debug: vi.fn(),
    info: vi.fn(),
    warn: vi.fn(),
    error: vi.fn(),
  },
}));

const importModule = vi.mocked(esmModule.importModule);
const runPython = vi.mocked(pythonUtils.runPython);
const readFileMock = vi.mocked(fs.promises.readFile);

describe('fileReference utility functions', () => {
  beforeEach(() => {
    vi.resetAllMocks();

    vi.mocked(logger.debug).mockImplementation((_message: string) => {
      return {
        debug: vi.fn(),
        info: vi.fn(),
        warn: vi.fn(),
        error: vi.fn(),
      } as unknown as Logger;
    });

    vi.mocked(logger.error).mockImplementation((_message: string) => {
      return {
        debug: vi.fn(),
        info: vi.fn(),
        warn: vi.fn(),
        error: vi.fn(),
      } as unknown as Logger;
    });

    vi.mocked(path.resolve).mockImplementation((basePath, filePath) =>
      filePath?.startsWith('/') ? filePath : path.join(basePath || '', filePath || ''),
    );

    vi.mocked(path.join).mockImplementation((...parts) => parts.filter(Boolean).join('/'));

    vi.mocked(path.extname).mockImplementation((filePath) => {
      if (!filePath) {
        return '';
      }
      const parts = filePath.split('.');
      return parts.length > 1 ? `.${parts[parts.length - 1]}` : '';
    });

    vi.mocked(isJavascriptFile).mockImplementation((filePath) => {
      if (!filePath) {
        return false;
      }
      return ['.js', '.mjs', '.ts', '.cjs'].some((ext) => filePath.endsWith(ext));
    });

    importModule.mockResolvedValue({});
    runPython.mockResolvedValue({});
  });

  describe('loadFileReference', () => {
    it('should load JSON files correctly', async () => {
      const fileRef = 'file:///path/to/config.json';
      const fileContent = '{"name": "test", "value": 42}';
      const parsedContent = { name: 'test', value: 42 };

      readFileMock.mockResolvedValue(fileContent);
      vi.spyOn(JSON, 'parse').mockReturnValue(parsedContent);

      const result = await loadFileReference(fileRef);

      expect(fs.promises.readFile).toHaveBeenCalledWith('/path/to/config.json', 'utf8');
      expect(result).toEqual(parsedContent);
    });

    it('should load YAML files correctly', async () => {
      const fileRef = 'file:///path/to/config.yaml';
      const fileContent = 'name: test\nvalue: 42';
      const parsedContent = { name: 'test', value: 42 };

      readFileMock.mockResolvedValue(fileContent);
      vi.mocked(loadYaml).mockReturnValue(parsedContent);

      const result = await loadFileReference(fileRef);

      expect(fs.promises.readFile).toHaveBeenCalledWith('/path/to/config.yaml', 'utf8');
      expect(loadYaml).toHaveBeenCalledWith(fileContent);
      expect(result).toEqual(parsedContent);
    });

    it('should load JavaScript files correctly', async () => {
      const fileRef = 'file:///path/to/config.js';
      const moduleOutput = { settings: { temperature: 0.7 } };

      vi.mocked(isJavascriptFile).mockReturnValue(true);
      importModule.mockResolvedValue(moduleOutput);

      const result = await loadFileReference(fileRef);

      expect(importModule).toHaveBeenCalledWith('/path/to/config.js', undefined);
      expect(result).toEqual(moduleOutput);
    });

    it('should load JavaScript files with function name correctly', async () => {
      const fileRef = 'file:///path/to/config.js:getConfig';
      const moduleOutput = { getConfig: 'success' };

      vi.mocked(isJavascriptFile).mockReturnValue(true);
      importModule.mockResolvedValue(moduleOutput);

      const result = await loadFileReference(fileRef);

      expect(importModule).toHaveBeenCalledWith('/path/to/config.js', 'getConfig');
      expect(result).toEqual(moduleOutput);
    });

    it('should load Python files correctly', async () => {
      const fileRef = 'file:///path/to/config.py';
      const pythonOutput = { message: 'Hello from Python' };

      runPython.mockResolvedValue(pythonOutput);

      const result = await loadFileReference(fileRef);

      expect(runPython).toHaveBeenCalledWith('/path/to/config.py', 'get_config', []);
      expect(result).toEqual(pythonOutput);
    });

    it('should load Python files with function name correctly', async () => {
      const fileRef = 'file:///path/to/config.py:custom_func';
      const pythonOutput = { custom: true };

      runPython.mockResolvedValue(pythonOutput);

      const result = await loadFileReference(fileRef);

      expect(runPython).toHaveBeenCalledWith('/path/to/config.py', 'custom_func', []);
      expect(result).toEqual(pythonOutput);
    });

    it('should load text files correctly', async () => {
      const fileRef = 'file:///path/to/config.txt';
      const fileContent = 'This is a text file';

      readFileMock.mockResolvedValue(fileContent);

      const result = await loadFileReference(fileRef);

      expect(fs.promises.readFile).toHaveBeenCalledWith('/path/to/config.txt', 'utf8');
      expect(result).toEqual(fileContent);
    });

    it('should resolve file paths relative to the basePath', async () => {
      const fileRef = 'file://config.json';
      const basePath = '/base/path';
      const fileContent = '{"name": "test"}';
      const parsedContent = { name: 'test' };

      vi.mocked(path.resolve).mockReturnValue('/base/path/config.json');
      readFileMock.mockResolvedValue(fileContent);
      vi.spyOn(JSON, 'parse').mockReturnValue(parsedContent);

      const result = await loadFileReference(fileRef, basePath);

      expect(path.resolve).toHaveBeenCalledWith('/base/path', 'config.json');
      expect(fs.promises.readFile).toHaveBeenCalledWith('/base/path/config.json', 'utf8');
      expect(result).toEqual(parsedContent);
    });

    it('should throw an error for unsupported file types', async () => {
      const fileRef = 'file:///path/to/file.xyz';

      vi.mocked(path.extname).mockReturnValue('.xyz');
      vi.mocked(isJavascriptFile).mockReturnValue(false);

      await expect(loadFileReference(fileRef)).rejects.toThrow('Unsupported file extension: .xyz');
    });
  });

  describe('processConfigFileReferences', () => {
    it('should return primitive values as is', async () => {
      await expect(processConfigFileReferences(42)).resolves.toBe(42);
      await expect(processConfigFileReferences('test')).resolves.toBe('test');
      await expect(processConfigFileReferences(true)).resolves.toBe(true);
      await expect(processConfigFileReferences(null)).resolves.toBeNull();
      await expect(processConfigFileReferences(undefined)).resolves.toBeUndefined();
    });

    it('should process a simple file reference string', async () => {
      const config = 'file:///path/to/config.json';
      const parsedContent = { name: 'test', value: 42 };

      readFileMock.mockResolvedValue(JSON.stringify(parsedContent));
      vi.spyOn(JSON, 'parse').mockReturnValue(parsedContent);

      const result = await processConfigFileReferences(config);

      expect(result).toEqual(parsedContent);
    });

    it('should process nested file references in objects', async () => {
      const config = {
        setting1: 'value1',
        setting2: 'file:///path/to/setting2.json',
        nested: {
          setting3: 'file:///path/to/setting3.yaml',
        },
      };

      readFileMock.mockImplementation((path) => {
        if (path === '/path/to/setting2.json') {
          return Promise.resolve('{"key": "value2"}');
        }
        if (path === '/path/to/setting3.yaml') {
          return Promise.resolve('key: value3');
        }
        return Promise.resolve('');
      });

      vi.spyOn(JSON, 'parse').mockImplementation((content) => {
        if (content === '{"key": "value2"}') {
          return { key: 'value2' };
        }
        return {};
      });

      vi.mocked(loadYaml).mockImplementation((content) => {
        if (content === 'key: value3') {
          return { key: 'value3' };
        }
        return {};
      });

      const result = await processConfigFileReferences(config);

      expect(result).toEqual({
        setting1: 'value1',
        setting2: { key: 'value2' },
        nested: {
          setting3: { key: 'value3' },
        },
      });
    });

    it('should process file references in arrays', async () => {
      const config = [
        'regular string',
        'file:///path/to/item1.json',
        'file:///path/to/item2.yaml',
        42,
      ];

      readFileMock.mockImplementation((path) => {
        if (path === '/path/to/item1.json') {
          return Promise.resolve('{"name": "item1"}');
        }
        if (path === '/path/to/item2.yaml') {
          return Promise.resolve('name: item2');
        }
        return Promise.resolve('');
      });

      vi.spyOn(JSON, 'parse').mockImplementation((content) => {
        if (content === '{"name": "item1"}') {
          return { name: 'item1' };
        }
        return {};
      });

      vi.mocked(loadYaml).mockImplementation((content) => {
        if (content === 'name: item2') {
          return { name: 'item2' };
        }
        return {};
      });

      const result = await processConfigFileReferences(config);

      expect(result).toEqual(['regular string', { name: 'item1' }, { name: 'item2' }, 42]);
    });

    it('should handle errors when processing file references', async () => {
      const config = {
        valid: 'regular value',
        invalid: 'file:///path/to/nonexistent.json',
      };

      readFileMock.mockImplementation((path) => {
        if (path === '/path/to/nonexistent.json') {
          return Promise.reject(new Error('File not found'));
        }
        return Promise.resolve('');
      });

      await expect(processConfigFileReferences(config)).rejects.toThrow('File not found');
      expect(logger.error).toHaveBeenCalledWith(expect.stringContaining('File not found'));
    });
  });
});