promptfoo--promptfoo/test/python/pythonUtils.test.ts

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

import { beforeEach, describe, expect, it, vi } from 'vitest';

// Create mock for execFileAsync - must be hoisted for vi.mock factory
const { mockExecFileAsync, mockExecFile } = vi.hoisted(() => {
  const mockExecFileAsync = vi.fn();
  // Create a mock execFile with the custom promisify symbol
  const mockExecFile = Object.assign(vi.fn(), {
    [Symbol.for('nodejs.util.promisify.custom')]: mockExecFileAsync,
  });
  return { mockExecFileAsync, mockExecFile };
});

// Mock child_process.execFile with custom promisify support
vi.mock('child_process', () => ({
  execFile: mockExecFile,
}));

import { PythonShell } from 'python-shell';
import { getEnvBool, getEnvString } from '../../src/envars';
import logger from '../../src/logger';
import * as pythonUtils from '../../src/python/pythonUtils';
import {
  createSecureTempDirectory,
  removeSecureTempDirectory,
  writeSecureTempFile,
} from '../../src/util/secureTempFiles';

const fsMock = vi.hoisted(() => ({
  writeFileSync: vi.fn(),
  readFileSync: vi.fn(),
  unlinkSync: vi.fn(),
}));

// Mock setup
vi.mock('fs', () => {
  return {
    ...fsMock,
    default: fsMock,
  };
});

vi.mock('fs/promises', () => ({
  default: {
    writeFile: fsMock.writeFileSync,
    readFile: fsMock.readFileSync,
    unlink: fsMock.unlinkSync,
  },
  writeFile: fsMock.writeFileSync,
  readFile: fsMock.readFileSync,
  unlink: fsMock.unlinkSync,
}));

vi.mock('../../src/envars', () => ({
  getEnvString: vi.fn(),
  getEnvBool: vi.fn(),
}));

vi.mock('../../src/logger', () => ({
  default: {
    debug: vi.fn(),
    error: vi.fn(),
    info: vi.fn(),
    warn: vi.fn(),
  },
}));

vi.mock('../../src/util/secureTempFiles', () => ({
  createSecureTempDirectory: vi.fn(),
  removeSecureTempDirectory: vi.fn(),
  writeSecureTempFile: vi.fn(),
}));

// Must be hoisted for vi.mock factory
const { mockPythonShellInstance, MockPythonShell } = vi.hoisted(() => {
  const instance = {
    stdout: { on: vi.fn() },
    stderr: { on: vi.fn() },
    end: vi.fn(),
  };
  // Create a proper class that can be used with 'new'
  const MockPythonShell = vi.fn(function (this: typeof instance) {
    Object.assign(this, instance);
    return this;
  }) as unknown as typeof import('python-shell').PythonShell;
  return { mockPythonShellInstance: instance, MockPythonShell };
});

vi.mock('python-shell', () => ({
  PythonShell: MockPythonShell,
}));

describe('Python Utils', () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mockExecFileAsync.mockReset();
    pythonUtils.state.cachedPythonPath = null;
    pythonUtils.state.validationPromise = null;
    mockPythonShellInstance.stdout.on.mockReset();
    mockPythonShellInstance.stderr.on.mockReset();
    mockPythonShellInstance.end.mockReset();
    // Set default mock return values
    vi.mocked(getEnvString).mockReturnValue('');
    vi.mocked(getEnvBool).mockReturnValue(false);
    vi.mocked(createSecureTempDirectory).mockResolvedValue('/tmp/promptfoo-python-test');
    vi.mocked(removeSecureTempDirectory).mockResolvedValue(undefined);
    vi.mocked(writeSecureTempFile).mockImplementation(
      async (directory: string, filename: string) => `${directory}/${filename}`,
    );
  });

  describe('getConfiguredPythonPath', () => {
    it('should return explicit config path when provided', () => {
      const result = pythonUtils.getConfiguredPythonPath('/custom/python/path');
      expect(result).toBe('/custom/python/path');
    });

    it('should return PROMPTFOO_PYTHON when config path is not provided', () => {
      vi.mocked(getEnvString).mockReturnValue('/env/python/path');

      const result = pythonUtils.getConfiguredPythonPath(undefined);

      expect(result).toBe('/env/python/path');
      expect(getEnvString).toHaveBeenCalledWith('PROMPTFOO_PYTHON');
    });

    it('should prioritize config path over PROMPTFOO_PYTHON', () => {
      vi.mocked(getEnvString).mockReturnValue('/env/python/path');

      const result = pythonUtils.getConfiguredPythonPath('/config/python/path');

      expect(result).toBe('/config/python/path');
    });

    it('should return undefined when neither config nor env var is set', () => {
      vi.mocked(getEnvString).mockReturnValue('');

      const result = pythonUtils.getConfiguredPythonPath(undefined);

      expect(result).toBeUndefined();
    });

    it('should return undefined when config is empty string and env var is not set', () => {
      vi.mocked(getEnvString).mockReturnValue('');

      const result = pythonUtils.getConfiguredPythonPath('');

      expect(result).toBeUndefined();
    });

    it('should return PROMPTFOO_PYTHON when config is empty string', () => {
      vi.mocked(getEnvString).mockReturnValue('/env/python');

      const result = pythonUtils.getConfiguredPythonPath('');

      expect(result).toBe('/env/python');
    });
  });

  describe('getSysExecutable', () => {
    it('should return Python executable path from sys.executable', async () => {
      // Mock for Unix-like systems (not Windows)
      const originalPlatform = process.platform;
      Object.defineProperty(process, 'platform', { value: 'linux' });

      mockExecFileAsync.mockResolvedValue({ stdout: '/usr/bin/python3.8\n', stderr: '' });

      const result = await pythonUtils.getSysExecutable();

      expect(result).toBe('/usr/bin/python3.8');
      expect(mockExecFileAsync).toHaveBeenCalledWith('python3', [
        '-c',
        'import sys; print(sys.executable)',
      ]);

      // Restore original platform
      Object.defineProperty(process, 'platform', { value: originalPlatform });
    });

    it('should use Windows where command first on Windows', async () => {
      const originalPlatform = process.platform;
      Object.defineProperty(process, 'platform', { value: 'win32' });

      let callCount = 0;
      mockExecFileAsync.mockImplementation(async () => {
        callCount++;
        if (callCount === 1) {
          // 'where python' returns multiple paths
          return {
            stdout:
              'C:\\Users\\test\\AppData\\Local\\Microsoft\\WindowsApps\\python.exe\nC:\\Python39\\python.exe\n',
            stderr: '',
          };
        } else {
          // Version check succeeds
          return { stdout: 'Python 3.9.0\n', stderr: '' };
        }
      });

      const result = await pythonUtils.getSysExecutable();

      // Should skip WindowsApps and use the real Python installation
      expect(result).toBe('C:\\Python39\\python.exe');
      expect(mockExecFileAsync).toHaveBeenCalledWith('where', ['python']);
      // Verify that the non-WindowsApps path was validated
      expect(mockExecFileAsync).toHaveBeenCalledWith('C:\\Python39\\python.exe', ['--version']);

      Object.defineProperty(process, 'platform', { value: originalPlatform });
    });

    it('should fall back to py commands if Windows where fails', async () => {
      const originalPlatform = process.platform;
      Object.defineProperty(process, 'platform', { value: 'win32' });

      let callCount = 0;
      mockExecFileAsync.mockImplementation(async () => {
        callCount++;
        if (callCount === 1) {
          // 'where python' fails
          throw new Error('where command failed');
        } else {
          // 'py' sys.executable succeeds
          return { stdout: 'C:\\Python39\\python.exe\n', stderr: '' };
        }
      });

      const result = await pythonUtils.getSysExecutable();

      expect(result).toBe('C:\\Python39\\python.exe');
      expect(mockExecFileAsync).toHaveBeenCalledWith('where', ['python']);
      // Verify py launcher fallback was used
      expect(mockExecFileAsync).toHaveBeenCalledWith('py', [
        '-c',
        'import sys; print(sys.executable)',
      ]);

      Object.defineProperty(process, 'platform', { value: originalPlatform });
    });

    it('should try direct python command as final Windows fallback', async () => {
      const originalPlatform = process.platform;
      Object.defineProperty(process, 'platform', { value: 'win32' });

      let callCount = 0;
      mockExecFileAsync.mockImplementation(async () => {
        callCount++;
        if (callCount <= 3) {
          // First 3 calls fail
          throw new Error('failed');
        } else {
          // Final fallback succeeds
          return { stdout: 'Python 3.9.0\n', stderr: '' };
        }
      });

      const result = await pythonUtils.getSysExecutable();

      expect(result).toBe('python');
      expect(mockExecFileAsync).toHaveBeenCalledWith('where', ['python']);
      // Verify the final fallback python --version was called
      expect(mockExecFileAsync).toHaveBeenCalledWith('python', ['--version']);

      Object.defineProperty(process, 'platform', { value: originalPlatform });
    });

    it('should add .exe suffix on Windows if missing', async () => {
      const originalPlatform = process.platform;
      Object.defineProperty(process, 'platform', { value: 'win32' });

      let callCount = 0;
      mockExecFileAsync.mockImplementation(async () => {
        callCount++;
        if (callCount === 1) {
          throw new Error('where failed');
        } else {
          return { stdout: 'C:\\Python39\\python\n', stderr: '' };
        }
      });

      const result = await pythonUtils.getSysExecutable();

      expect(result).toBe('C:\\Python39\\python.exe');

      Object.defineProperty(process, 'platform', { value: originalPlatform });
    });

    it('should handle empty where output gracefully', async () => {
      const originalPlatform = process.platform;
      Object.defineProperty(process, 'platform', { value: 'win32' });

      let callCount = 0;
      mockExecFileAsync.mockImplementation(async () => {
        callCount++;
        if (callCount === 1) {
          // 'where python' returns empty
          return { stdout: '', stderr: '' };
        } else {
          // 'py' succeeds
          return { stdout: 'C:\\Python39\\python.exe\n', stderr: '' };
        }
      });

      const result = await pythonUtils.getSysExecutable();

      expect(result).toBe('C:\\Python39\\python.exe');
      expect(mockExecFileAsync).toHaveBeenCalledWith('where', ['python']);
      // Verify py launcher fallback was used when where returned empty
      expect(mockExecFileAsync).toHaveBeenCalledWith('py', [
        '-c',
        'import sys; print(sys.executable)',
      ]);

      Object.defineProperty(process, 'platform', { value: originalPlatform });
    });

    it('should return null if no Python executable is found', async () => {
      mockExecFileAsync.mockRejectedValue(new Error('Command failed'));

      const result = await pythonUtils.getSysExecutable();

      expect(result).toBeNull();
    });
  });

  describe('tryPath', () => {
    describe('successful path validation', () => {
      it('should return the path for a valid Python 3 executable', async () => {
        mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

        const result = await pythonUtils.tryPath('/usr/bin/python3');

        expect(result).toBe('/usr/bin/python3');
        expect(mockExecFileAsync).toHaveBeenCalledWith('/usr/bin/python3', ['--version']);
      });
    });

    describe('failed path validation', () => {
      it('should return null for a non-existent executable', async () => {
        mockExecFileAsync.mockRejectedValue(new Error('Command failed'));

        const result = await pythonUtils.tryPath('/usr/bin/nonexistent');

        expect(result).toBeNull();
        expect(mockExecFileAsync).toHaveBeenCalledWith('/usr/bin/nonexistent', ['--version']);
      });

      it('should return null if the command times out', async () => {
        vi.useFakeTimers();

        // Mock execFileAsync to return a promise that never resolves (simulating timeout)
        mockExecFileAsync.mockImplementation(() => new Promise(() => {}));

        const resultPromise = pythonUtils.tryPath('/usr/bin/python3');
        await vi.advanceTimersByTimeAsync(2501);

        const result = await resultPromise;

        expect(result).toBeNull();
        expect(mockExecFileAsync).toHaveBeenCalledWith('/usr/bin/python3', ['--version']);
        vi.useRealTimers();
      });
    });
  });

  describe('validatePythonPath', () => {
    describe('caching behavior', () => {
      it('should validate and cache an existing Python 3 path', async () => {
        mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

        const result = await pythonUtils.validatePythonPath('python', false);

        expect(result).toBe('python');
        expect(pythonUtils.state.cachedPythonPath).toBe('python');
        expect(mockExecFileAsync).toHaveBeenCalledWith('python', ['--version']);
      });

      it('should return the cached path on subsequent calls', async () => {
        pythonUtils.state.cachedPythonPath = '/usr/bin/python3';

        const result = await pythonUtils.validatePythonPath('python', false);

        expect(result).toBe('/usr/bin/python3');
        expect(mockExecFileAsync).not.toHaveBeenCalled();
      });
    });

    describe('fallback behavior', () => {
      it('should fall back to alternative paths for non-existent programs when not explicit', async () => {
        mockExecFileAsync.mockReset();
        let callCount = 0;

        mockExecFileAsync.mockImplementation(async () => {
          callCount++;
          if (callCount === 1) {
            // Primary path fails
            throw new Error('Command failed');
          } else if (process.platform === 'win32') {
            if (callCount <= 4) {
              throw new Error('Command failed');
            } else {
              return { stdout: 'Python 3.9.5\n', stderr: '' };
            }
          } else {
            // Unix: getSysExecutable succeeds on python3
            return { stdout: '/usr/bin/python3\n', stderr: '' };
          }
        });

        const result = await pythonUtils.validatePythonPath('non_existent_program', false);

        expect(result).toBe(process.platform === 'win32' ? 'python' : '/usr/bin/python3');
      });

      it('should throw an error for non-existent programs when explicit', async () => {
        mockExecFileAsync.mockRejectedValue(new Error('Command failed'));

        await expect(pythonUtils.validatePythonPath('non_existent_program', true)).rejects.toThrow(
          'Python 3 not found. Tried "non_existent_program"',
        );
        expect(mockExecFileAsync).toHaveBeenCalledWith('non_existent_program', ['--version']);
      });

      it('should throw an error when no valid Python path is found', async () => {
        mockExecFileAsync.mockReset();
        mockExecFileAsync.mockRejectedValue(new Error('Command failed'));

        await expect(pythonUtils.validatePythonPath('python', false)).rejects.toThrow(
          'Python 3 not found. Tried "python", sys.executable detection, and fallback commands.',
        );
        expect(mockExecFileAsync).toHaveBeenCalled();
      });
    });

    describe('environment variable handling', () => {
      it('should use PROMPTFOO_PYTHON environment variable when provided', async () => {
        vi.mocked(getEnvString).mockReturnValue('/custom/python/path');
        mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

        const result = await pythonUtils.validatePythonPath('/custom/python/path', true);

        expect(result).toBe('/custom/python/path');
        expect(mockExecFileAsync).toHaveBeenCalledWith('/custom/python/path', ['--version']);
      });
    });

    describe('concurrent validation', () => {
      it('should share validation promise between concurrent calls', async () => {
        mockExecFileAsync.mockImplementation(async () => {
          // Yield control so concurrent callers can register on the shared
          // validation promise before resolution.
          await Promise.resolve();
          return { stdout: 'Python 3.8.10\n', stderr: '' };
        });

        // Start two validations concurrently
        const [result1, result2] = await Promise.all([
          pythonUtils.validatePythonPath('python', false),
          pythonUtils.validatePythonPath('python', false),
        ]);

        expect(result1).toBe('python');
        expect(result2).toBe('python');

        // Only one exec call should be made
        expect(mockExecFileAsync).toHaveBeenCalledTimes(1);

        // After resolution, validation promise should be cleared
        expect(pythonUtils.state.validationPromise).toBeNull();
      });

      it('should handle race conditions between concurrent validation attempts', async () => {
        // Clear cached path first to ensure validation runs
        pythonUtils.state.cachedPythonPath = null;

        mockExecFileAsync.mockImplementation(async () => {
          // Yield control so concurrent callers can race onto the shared
          // validation promise before resolution.
          await Promise.resolve();
          return { stdout: 'Python 3.8.10\n', stderr: '' };
        });

        // Start multiple validations without waiting
        const promises = [
          pythonUtils.validatePythonPath('python', false),
          pythonUtils.validatePythonPath('python', false),
          pythonUtils.validatePythonPath('python', false),
        ];

        const results = await Promise.all(promises);

        expect(results).toEqual(['python', 'python', 'python']);

        // Only one exec call should be made
        expect(mockExecFileAsync).toHaveBeenCalledTimes(1);

        // After resolution, validation promise should be cleared
        expect(pythonUtils.state.validationPromise).toBeNull();
      });
    });

    describe('promise cleanup', () => {
      it('should clear validation promise after failed validation', async () => {
        mockExecFileAsync.mockRejectedValue(new Error('Command failed'));

        await expect(pythonUtils.validatePythonPath('python', true)).rejects.toThrow(
          'Python 3 not found. Tried "python"',
        );

        // Validation promise should be cleared even after failure
        expect(pythonUtils.state.validationPromise).toBeNull();
      });
    });
  });

  describe('runPython', () => {
    beforeEach(() => {
      vi.clearAllMocks();
    });

    it('should execute a Python script with proper arguments', async () => {
      vi.mocked(fs.writeFileSync).mockImplementation(() => {});
      vi.mocked(fs.readFileSync).mockReturnValue(
        JSON.stringify({ type: 'final_result', data: 42 }),
      );
      vi.mocked(fs.unlinkSync).mockImplementation(() => {});
      mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

      mockPythonShellInstance.end.mockImplementation((callback: any) => {
        callback(null);
      });

      const scriptPath = '/path/to/script.py';
      const result = await pythonUtils.runPython(scriptPath, 'test_method', [1, 2, 3]);

      expect(result).toBe(42);
      expect(createSecureTempDirectory).toHaveBeenCalledWith('promptfoo-python-');
      expect(writeSecureTempFile).toHaveBeenNthCalledWith(
        1,
        '/tmp/promptfoo-python-test',
        'input.json',
        '[1,2,3]',
      );
      expect(writeSecureTempFile).toHaveBeenNthCalledWith(
        2,
        '/tmp/promptfoo-python-test',
        'output.json',
        '',
      );
      expect(removeSecureTempDirectory).toHaveBeenCalledWith('/tmp/promptfoo-python-test');
      expect(PythonShell).toHaveBeenCalledWith(
        'wrapper.py',
        expect.objectContaining({
          scriptPath: expect.any(String),
          args: expect.arrayContaining([path.resolve(scriptPath), 'test_method']),
        }),
      );
    });

    it('should use python subdirectory for wrapper.py scriptPath', async () => {
      vi.mocked(fs.writeFileSync).mockImplementation(() => {});
      vi.mocked(fs.readFileSync).mockReturnValue(
        JSON.stringify({ type: 'final_result', data: 'test_result' }),
      );
      vi.mocked(fs.unlinkSync).mockImplementation(() => {});
      mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

      mockPythonShellInstance.end.mockImplementation((callback: any) => {
        callback(null);
      });

      const scriptPath = '/path/to/script.py';
      await pythonUtils.runPython(scriptPath, 'test_method', ['arg1', 'arg2']);

      // Verify PythonShell was called with scriptPath ending in 'python'
      // This works for both development (src/python/) and production (dist/src/python/)
      expect(PythonShell).toHaveBeenCalledWith(
        'wrapper.py',
        expect.objectContaining({
          scriptPath: expect.stringMatching(/python$/),
        }),
      );
    });

    it('should throw an error if Python script returns invalid JSON', async () => {
      vi.mocked(fs.writeFileSync).mockImplementation(() => {});
      vi.mocked(fs.readFileSync).mockReturnValue('invalid json');
      vi.mocked(fs.unlinkSync).mockImplementation(() => {});
      mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

      mockPythonShellInstance.end.mockImplementation((callback: any) => {
        callback(null);
      });

      await expect(pythonUtils.runPython('/path/to/script.py', 'test_method', [])).rejects.toThrow(
        'Invalid JSON',
      );
    });

    it('should not log arguments or returned payloads', async () => {
      vi.mocked(fs.readFileSync).mockReturnValue(
        JSON.stringify({ type: 'final_result', data: 'secret-result' }),
      );
      mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

      mockPythonShellInstance.end.mockImplementation((callback: any) => {
        callback(null);
      });

      await pythonUtils.runPython('/path/to/script.py', 'test_method', ['secret-input']);

      const debugMessages = vi.mocked(logger.debug).mock.calls.flat().map(String).join('\n');
      expect(debugMessages).not.toContain('secret-input');
      expect(debugMessages).not.toContain('secret-result');
    });

    it('classifies routine stderr without reporting Python logging as errors', async () => {
      vi.mocked(fs.readFileSync).mockReturnValue(
        JSON.stringify({ type: 'final_result', data: 'result' }),
      );
      mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });
      mockPythonShellInstance.stderr.on.mockImplementation((event: string, callback: any) => {
        if (event === 'data') {
          callback(Buffer.from('INFO:root:loaded config\nWARNING:root:slow response\n'));
        }
      });
      mockPythonShellInstance.end.mockImplementation((callback: any) => {
        callback(null);
      });

      await pythonUtils.runPython('/path/to/script.py', 'test_method', []);

      expect(logger.info).toHaveBeenCalledWith('INFO:root:loaded config');
      expect(logger.warn).toHaveBeenCalledWith('WARNING:root:slow response');
      expect(logger.error).not.toHaveBeenCalled();
    });

    it('should throw an error if Python script does not return final_result', async () => {
      vi.mocked(fs.writeFileSync).mockImplementation(() => {});
      vi.mocked(fs.readFileSync).mockReturnValue(JSON.stringify({ type: 'other', data: 42 }));
      vi.mocked(fs.unlinkSync).mockImplementation(() => {});
      mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

      mockPythonShellInstance.end.mockImplementation((callback: any) => {
        callback(null);
      });

      await expect(pythonUtils.runPython('/path/to/script.py', 'test_method', [])).rejects.toThrow(
        'The Python script `call_api` function must return a dict with an `output`',
      );
    });

    it('should clean up temporary files even on error', async () => {
      vi.mocked(fs.writeFileSync).mockImplementation(() => {});
      vi.mocked(fs.readFileSync).mockImplementation(() => {
        throw new Error('Read failed');
      });
      vi.mocked(fs.unlinkSync).mockImplementation(() => {});
      mockExecFileAsync.mockResolvedValue({ stdout: 'Python 3.8.10\n', stderr: '' });

      mockPythonShellInstance.end.mockImplementation((callback: any) => {
        callback(null);
      });

      await expect(
        pythonUtils.runPython('/path/to/script.py', 'test_method', []),
      ).rejects.toThrow();

      expect(removeSecureTempDirectory).toHaveBeenCalledWith('/tmp/promptfoo-python-test');
    });
  });
});