From cdb361b5a3552a0951faa5e028757509f8b71719 Mon Sep 17 00:00:00 2001 From: mohammed ahmed Date: Fri, 3 Apr 2026 13:44:27 +0000 Subject: [PATCH] Fix vi.mock() path resolution in generated vitest tests Extended fix_jest_mock_paths() to handle vitest mock calls (vi.mock()) in addition to jest.mock(). Previously, only jest.mock() paths were corrected, causing vitest tests to fail with "Cannot find module" errors. Problem: - Source at src/agents/workspace.ts imports ../routing/session-key - Generated test at test/test_workspace.test.ts used vi.mock('../routing/session-key') - This resolves to /routing/session-key (wrong - goes up from test/, not found) - Should be vi.mock('../src/routing/session-key') (correct path from test/) Solution: - Updated regex pattern to match both jest.mock() and vi.mock() - Function now fixes relative paths for both test frameworks - Added unit tests to verify both jest and vitest paths are corrected Trace ID: 265059d4-f518-44da-8367-d90ca424092c Co-Authored-By: Claude Sonnet 4.5 --- codeflash/languages/javascript/instrument.py | 14 +-- tests/test_fix_mock_paths_vitest.py | 94 ++++++++++++++++++++ 2 files changed, 101 insertions(+), 7 deletions(-) create mode 100644 tests/test_fix_mock_paths_vitest.py diff --git a/codeflash/languages/javascript/instrument.py b/codeflash/languages/javascript/instrument.py index 8bcd0b2ee..cfce9b224 100644 --- a/codeflash/languages/javascript/instrument.py +++ b/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): diff --git a/tests/test_fix_mock_paths_vitest.py b/tests/test_fix_mock_paths_vitest.py new file mode 100644 index 000000000..b37bbda7d --- /dev/null +++ b/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!")