chore: import upstream snapshot with attribution

test/python/wrapper.test.ts

@@ -0,0 +1,213 @@
+import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import {
+  getSysExecutable,
+  runPython,
+  state,
+  tryPath,
+  validatePythonPath,
+} from '../../src/python/pythonUtils';
+import { runPythonCode } from '../../src/python/wrapper';
+import {
+  createSecureTempDirectory,
+  removeSecureTempDirectory,
+  writeSecureTempFile,
+} from '../../src/util/secureTempFiles';
+import { mockProcessEnv } from '../util/utils';
+
+vi.mock('../../src/esm');
+vi.mock('python-shell');
+vi.mock('../../src/util/secureTempFiles', () => ({
+  createSecureTempDirectory: vi.fn(),
+  removeSecureTempDirectory: vi.fn(),
+  writeSecureTempFile: vi.fn(),
+}));
+vi.mock('../../src/python/pythonUtils', async () => {
+  const originalModule = await vi.importActual<typeof import('../../src/python/pythonUtils')>(
+    '../../src/python/pythonUtils',
+  );
+  return {
+    ...originalModule,
+    validatePythonPath: vi.fn(),
+    runPython: vi.fn(originalModule.runPython),
+    tryPath: vi.fn(),
+    getSysExecutable: vi.fn(),
+    // Use the real state object so all implementations share the same cache
+    state: originalModule.state,
+  };
+});
+interface TestResult {
+  testId: number;
+  result: string;
+}
+interface MixedTestResult {
+  testId: number;
+  result: string;
+  isExplicit: boolean;
+}
+describe('wrapper', () => {
+  let restoreEnv: () => void;
+
+  beforeAll(() => {
+    restoreEnv = mockProcessEnv({ PROMPTFOO_PYTHON: undefined });
+  });
+
+  afterAll(() => {
+    restoreEnv();
+  });
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.mocked(createSecureTempDirectory).mockResolvedValue('/tmp/promptfoo-python-code-test');
+    vi.mocked(removeSecureTempDirectory).mockResolvedValue(undefined);
+    vi.mocked(writeSecureTempFile).mockImplementation(
+      async (directory: string, filename: string) => `${directory}/${filename}`,
+    );
+    vi.mocked(validatePythonPath).mockImplementation(
+      (pythonPath: string, _isExplicit: boolean): Promise<string> => {
+        state.cachedPythonPath = pythonPath;
+        return Promise.resolve(pythonPath);
+      },
+    );
+  });
+  describe('runPythonCode', () => {
+    it('should clean up the temporary files after execution', async () => {
+      const mockRunPython = vi.fn().mockResolvedValue('cleanup test');
+      vi.mocked(runPython).mockImplementation(mockRunPython);
+      await runPythonCode('print("cleanup test")', 'main', []);
+      expect(createSecureTempDirectory).toHaveBeenCalledWith('promptfoo-python-code-');
+      expect(writeSecureTempFile).toHaveBeenCalledWith(
+        '/tmp/promptfoo-python-code-test',
+        'script.py',
+        'print("cleanup test")',
+      );
+      expect(mockRunPython).toHaveBeenCalledTimes(1);
+      expect(mockRunPython).toHaveBeenCalledWith(
+        '/tmp/promptfoo-python-code-test/script.py',
+        'main',
+        [],
+      );
+      expect(removeSecureTempDirectory).toHaveBeenCalledWith('/tmp/promptfoo-python-code-test');
+    });
+    it('should execute Python code from a string and read the output file', async () => {
+      const mockOutput = { type: 'final_result', data: 'execution result' };
+      const mockRunPython = vi.mocked(runPython);
+      mockRunPython.mockResolvedValue(mockOutput.data);
+      const code = 'print("Hello, world!")';
+      const result = await runPythonCode(code, 'main', []);
+      expect(result).toBe('execution result');
+      expect(mockRunPython).toHaveBeenCalledWith(expect.stringContaining('.py'), 'main', []);
+      expect(writeSecureTempFile).toHaveBeenCalledWith(
+        '/tmp/promptfoo-python-code-test',
+        'script.py',
+        code,
+      );
+    });
+  });
+  describe('validatePythonPath race conditions', () => {
+    beforeAll(async () => {
+      // Restore the real validatePythonPath implementation while keeping
+      // tryPath and getSysExecutable mocked to avoid actual system calls
+      const realModule = await vi.importActual<typeof import('../../src/python/pythonUtils')>(
+        '../../src/python/pythonUtils',
+      );
+      vi.mocked(validatePythonPath).mockImplementation(realModule.validatePythonPath);
+    });
+
+    beforeEach(() => {
+      // Reset the cached path and validation promise before each test
+      state.cachedPythonPath = null;
+      state.validationPromise = null;
+
+      // Yield once before resolving so concurrent callers can race onto the
+      // shared validation promise — this is the same race condition the test
+      // is exercising; wall-clock delay isn't needed.
+      vi.mocked(tryPath).mockImplementation(async (_path: string) => {
+        await Promise.resolve();
+        return '/usr/bin/python3';
+      });
+      vi.mocked(getSysExecutable).mockImplementation(async () => {
+        await Promise.resolve();
+        return '/usr/bin/python3';
+      });
+    });
+    it('should handle concurrent validatePythonPath calls without race conditions', async () => {
+      // Launch multiple concurrent validations
+      const concurrentCalls = 10;
+      const promises = Array.from({ length: concurrentCalls }, (_, i) =>
+        validatePythonPath('python', false).then(
+          (result: string): TestResult => ({ testId: i, result }),
+        ),
+      );
+      const results = await Promise.all(promises);
+      // All calls should succeed
+      expect(results).toHaveLength(concurrentCalls);
+      results.forEach((result: TestResult) => {
+        expect(result.result).toBeTruthy();
+        expect(typeof result.result).toBe('string');
+      });
+      // All results should be consistent (no race condition)
+      // If there was a race condition, different calls might get different results
+      const uniqueResults = new Set(results.map((r: TestResult) => r.result));
+      expect(uniqueResults.size).toBe(1);
+      // Cache should be populated
+      expect(state.cachedPythonPath).toBeTruthy();
+    }, 10000); // Increase timeout for this test
+    it('should handle mixed explicit/implicit validation calls consistently', async () => {
+      // Create mixed explicit/implicit calls
+      const mixedPromises = Array.from({ length: 8 }, (_, i) => {
+        const isExplicit = i % 2 === 0;
+        return validatePythonPath('python', isExplicit).then(
+          (result: string): MixedTestResult => ({
+            testId: i,
+            result,
+            isExplicit,
+          }),
+        );
+      });
+      const results = await Promise.all(mixedPromises);
+      // All calls should succeed
+      expect(results).toHaveLength(8);
+      results.forEach((result: MixedTestResult) => {
+        expect(result.result).toBeTruthy();
+        expect(typeof result.result).toBe('string');
+      });
+      // Check consistency between explicit and implicit results
+      const explicitResults = results
+        .filter((r: MixedTestResult) => r.isExplicit)
+        .map((r: MixedTestResult) => r.result);
+      const implicitResults = results
+        .filter((r: MixedTestResult) => !r.isExplicit)
+        .map((r: MixedTestResult) => r.result);
+      const uniqueExplicitResults = new Set(explicitResults);
+      const uniqueImplicitResults = new Set(implicitResults);
+      // Both explicit and implicit calls should return consistent results
+      // If there was a race condition, explicit and implicit calls might return different results
+      expect(uniqueExplicitResults.size).toBe(1);
+      expect(uniqueImplicitResults.size).toBe(1);
+      // Explicit and implicit results should be the same
+      expect(explicitResults[0]).toBe(implicitResults[0]);
+    }, 10000);
+    it('should handle rapid successive calls without race conditions', async () => {
+      // Launch rapid successive calls
+      const rapidCalls = 20;
+      const promises = Array.from({ length: rapidCalls }, (_, i) =>
+        validatePythonPath('python', false).then(
+          (result: string): TestResult => ({ testId: i, result }),
+        ),
+      );
+      const results = await Promise.all(promises);
+      // All calls should succeed
+      expect(results).toHaveLength(rapidCalls);
+      results.forEach((result: TestResult) => {
+        expect(result.result).toBeTruthy();
+        expect(typeof result.result).toBe('string');
+      });
+      // All results should be consistent
+      // If there was a race condition, different calls could get different cached values
+      const uniqueResults = new Set(results.map((r: TestResult) => r.result));
+      expect(uniqueResults.size).toBe(1);
+      // Cache should be populated with the consistent result
+      expect(state.cachedPythonPath).toBe(results[0].result);
+    }, 10000);
+  });
+});