Fix vi.mock() path resolution in generated vitest tests

codeflash/languages/javascript/instrument.py

@@ -1287,13 +1287,13 @@ def fix_imports_inside_test_blocks(test_code: str) -> str:
 
 
 def fix_jest_mock_paths(test_code: str, test_file_path: Path, source_file_path: Path, tests_root: Path) -> str:
-    """Fix relative paths in jest.mock() calls to be correct from the test file's location.
+    """Fix relative paths in jest.mock() and vi.mock() calls to be correct from the test file's location.
 
-    The AI sometimes generates jest.mock() calls with paths relative to the source file
+    The AI sometimes generates mock calls with paths relative to the source file
     instead of the test file. For example:
     - Source at `src/queue/queue.ts` imports `../environment` (-> src/environment)
-    - Test at `tests/test.test.ts` generates `jest.mock('../environment')` (-> ./environment, wrong!)
-    - Should generate `jest.mock('../src/environment')`
+    - Test at `tests/test.test.ts` generates `jest.mock('../environment')` or `vi.mock('../environment')` (-> ./environment, wrong!)
+    - Should generate `jest.mock('../src/environment')` or `vi.mock('../src/environment')`
 
     This function detects relative mock paths and adjusts them based on the test file's
     location relative to the source file's directory.
@@ -1318,8 +1318,8 @@ def fix_jest_mock_paths(test_code: str, test_file_path: Path, source_file_path:
     test_dir = test_file_path.resolve().parent
     project_root = tests_root.resolve().parent if tests_root.name == "tests" else tests_root.resolve()
 
-    # Pattern to match jest.mock() or jest.doMock() with relative paths
-    mock_pattern = re.compile(r"(jest\.(?:mock|doMock)\s*\(\s*['\"])(\.\./[^'\"]+|\.\/[^'\"]+)(['\"])")
+    # Pattern to match jest.mock(), jest.doMock(), or vi.mock() with relative paths
+    mock_pattern = re.compile(r"((?:jest|vi)\.(?:mock|doMock)\s*\(\s*['\"])(\.\./[^'\"]+|\.\/[^'\"]+)(['\"])")
 
     def fix_mock_path(match: re.Match[str]) -> str:
         original = match.group(0)
@@ -1359,7 +1359,7 @@ def fix_mock_path(match: re.Match[str]) -> str:
                 if not new_rel_path.startswith("../") and not new_rel_path.startswith("./"):
                     new_rel_path = f"./{new_rel_path}"
 
-                logger.debug(f"Fixed jest.mock path: {rel_path} -> {new_rel_path}")
+                logger.debug(f"Fixed mock path: {rel_path} -> {new_rel_path}")
                 return f"{prefix}{new_rel_path}{suffix}"
 
         except (ValueError, OSError):

tests/test_fix_mock_paths_vitest.py

@@ -0,0 +1,94 @@
+"""Test fix_jest_mock_paths function with vitest mocks."""
+
+from pathlib import Path
+
+from codeflash.languages.javascript.instrument import fix_jest_mock_paths
+
+
+def test_fix_vitest_mock_paths():
+    """Test that vi.mock() paths are fixed correctly."""
+    # Simulate source at src/agents/workspace.ts importing from ../routing/session-key
+    # Test at test/test_workspace.test.ts should mock ../src/routing/session-key, not ../routing/session-key
+
+    test_code = """
+vi.mock('../routing/session-key', () => ({
+  isSubagentSessionKey: vi.fn(),
+  isCronSessionKey: vi.fn(),
+}));
+
+import { filterBootstrapFilesForSession } from '../src/agents/workspace.js';
+    """
+
+    # Create temp directories and files for testing
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        project = Path(tmpdir)
+
+        # Create directory structure
+        src = project / "src"
+        src_agents = src / "agents"
+        src_routing = src / "routing"
+        test_dir = project / "test"
+
+        src_agents.mkdir(parents=True)
+        src_routing.mkdir(parents=True)
+        test_dir.mkdir(parents=True)
+
+        # Create files
+        source_file = src_agents / "workspace.ts"
+        source_file.write_text("export function filterBootstrapFilesForSession() {}")
+
+        routing_file = src_routing / "session-key.ts"
+        routing_file.write_text("export function isSubagentSessionKey() {}")
+
+        test_file = test_dir / "test_workspace.test.ts"
+        test_file.write_text(test_code)
+
+        # Fix the paths
+        fixed = fix_jest_mock_paths(test_code, test_file, source_file, test_dir)
+
+        # Should change ../routing/session-key to ../src/routing/session-key
+        assert "../src/routing/session-key" in fixed, f"Expected path to be fixed, got: {fixed}"
+        assert "../routing/session-key" not in fixed or "../src/routing/session-key" in fixed
+
+
+def test_fix_jest_mock_paths_still_works():
+    """Test that jest.mock() paths are still fixed correctly."""
+    test_code = """
+jest.mock('../routing/session-key', () => ({
+  isSubagentSessionKey: jest.fn(),
+}));
+    """
+
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        project = Path(tmpdir)
+        src = project / "src"
+        src_agents = src / "agents"
+        src_routing = src / "routing"
+        test_dir = project / "test"
+
+        src_agents.mkdir(parents=True)
+        src_routing.mkdir(parents=True)
+        test_dir.mkdir(parents=True)
+
+        source_file = src_agents / "workspace.ts"
+        source_file.write_text("")
+
+        routing_file = src_routing / "session-key.ts"
+        routing_file.write_text("")
+
+        test_file = test_dir / "test_workspace.test.ts"
+        test_file.write_text(test_code)
+
+        fixed = fix_jest_mock_paths(test_code, test_file, source_file, test_dir)
+
+        assert "../src/routing/session-key" in fixed
+
+
+if __name__ == "__main__":
+    test_fix_vitest_mock_paths()
+    test_fix_jest_mock_paths_still_works()
+    print("All tests passed!")