Merge branch 'main' into fix/vitest-coverage-override

Author: mohammedahmed18

codeflash/languages/function_optimizer.py

@@ -3085,6 +3085,16 @@ def run_optimized_candidate(
                 )
             )
 
+    def get_js_project_root(self) -> Path | None:
+        # Only calculate for JavaScript/TypeScript projects
+        if self.function_to_optimize.language not in ("javascript", "typescript"):
+            return self.test_cfg.js_project_root  # Fall back to cached value for non-JS
+
+        # For JS/TS, calculate fresh for each function to support monorepos
+        from codeflash.languages.javascript.test_runner import find_node_project_root
+
+        return find_node_project_root(Path(self.function_to_optimize.file_path))
+
     def run_and_parse_tests(
         self,
         testing_type: TestingMode,
@@ -3103,33 +3113,39 @@ def run_and_parse_tests(
         coverage_config_file = None
         try:
             if testing_type == TestingMode.BEHAVIOR:
+                # Calculate js_project_root for the current function being optimized
+                # instead of using cached value from test_cfg, which may be from a different function
+                js_project_root = self.get_js_project_root()
+
                 result_file_path, run_result, coverage_database_file, coverage_config_file = (
                     self.language_support.run_behavioral_tests(
                         test_paths=test_files,
                         test_env=test_env,
                         cwd=self.project_root,
                         timeout=INDIVIDUAL_TESTCASE_TIMEOUT,
-                        project_root=self.test_cfg.js_project_root,
+                        project_root=js_project_root,
                         enable_coverage=enable_coverage,
                         candidate_index=optimization_iteration,
                     )
                 )
             elif testing_type == TestingMode.LINE_PROFILE:
+                js_project_root = self.get_js_project_root()
                 result_file_path, run_result = self.language_support.run_line_profile_tests(
                     test_paths=test_files,
                     test_env=test_env,
                     cwd=self.project_root,
                     timeout=INDIVIDUAL_TESTCASE_TIMEOUT,
-                    project_root=self.test_cfg.js_project_root,
+                    project_root=js_project_root,
                     line_profile_output_file=line_profiler_output_file,
                 )
             elif testing_type == TestingMode.PERFORMANCE:
+                js_project_root = self.get_js_project_root()
                 result_file_path, run_result = self.language_support.run_benchmarking_tests(
                     test_paths=test_files,
                     test_env=test_env,
                     cwd=self.project_root,
                     timeout=INDIVIDUAL_TESTCASE_TIMEOUT,
-                    project_root=self.test_cfg.js_project_root,
+                    project_root=js_project_root,
                     min_loops=pytest_min_loops,
                     max_loops=pytest_max_loops,
                     target_duration_seconds=testing_time,

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):

codeflash/languages/javascript/vitest_runner.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import os
+import re
 import subprocess
 import time
 from pathlib import Path
@@ -169,7 +170,7 @@ def _is_vitest_workspace(project_root: Path) -> bool:
         return False
 
     try:
-        content = vitest_config.read_text()
+        content = vitest_config.read_text(encoding="utf-8")
         # Check for actual workspace configuration patterns (not just the word "workspace" in comments)
         # Valid indicators:
         #   - defineWorkspace() function call

codeflash/verification/verifier.py

@@ -34,7 +34,9 @@ def generate_tests(
     # TODO: Sometimes this recreates the original Class definition. This overrides and messes up the original
     #  class import. Remove the recreation of the class definition
     start_time = time.perf_counter()
-    test_module_path = Path(module_name_from_file_path(test_path, test_cfg.tests_project_rootdir))
+    # Use traverse_up=True to handle co-located __tests__ directories that may be outside
+    # the configured tests_root (e.g., src/gateway/__tests__/ when tests_root is test/)
+    test_module_path = Path(module_name_from_file_path(test_path, test_cfg.tests_project_rootdir, traverse_up=True))
 
     # Detect module system via language support (non-None for JS/TS, None for Python)
     lang_support = current_language_support()

tests/languages/javascript/test_vitest_setupfiles_fix.py

@@ -0,0 +1,61 @@
+from pathlib import Path
+
+import pytest
+
+from codeflash.languages.javascript.vitest_runner import _ensure_codeflash_vitest_config
+
+
+def test_codeflash_vitest_config_overrides_setupfiles(tmp_path: Path) -> None:
+    project_root = tmp_path.resolve()
+
+    # Create a project with setup file
+    (project_root / "test").mkdir()
+    (project_root / "test" / "setup.ts").write_text("// Setup file\n", encoding="utf-8")
+
+    vitest_config = """import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    setupFiles: ["test/setup.ts"],  // Relative path - will cause issues
+    include: ["src/**/*.test.ts"],
+  },
+});
+"""
+    (project_root / "vitest.config.ts").write_text(vitest_config, encoding="utf-8")
+
+    codeflash_config_path = _ensure_codeflash_vitest_config(project_root)
+
+    assert codeflash_config_path is not None
+    assert codeflash_config_path.exists()
+
+    config_content = codeflash_config_path.read_text(encoding="utf-8")
+
+    assert "setupFiles" in config_content, (
+        "Generated config must explicitly handle setupFiles to prevent "
+        "relative path resolution issues. Current config:\n" + config_content
+    )
+    assert "setupFiles: []" in config_content or "setupFiles:" in config_content, (
+        "setupFiles must be explicitly set in the merged config"
+    )
+
+
+def test_codeflash_vitest_config_without_setupfiles(tmp_path: Path) -> None:
+    project_root = tmp_path.resolve()
+
+    vitest_config = """import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    include: ["src/**/*.test.ts"],
+  },
+});
+"""
+    (project_root / "vitest.config.ts").write_text(vitest_config, encoding="utf-8")
+
+    codeflash_config_path = _ensure_codeflash_vitest_config(project_root)
+
+    assert codeflash_config_path is not None
+    assert codeflash_config_path.exists()
+
+    config_content = codeflash_config_path.read_text(encoding="utf-8")
+    assert "mergeConfig" in config_content or "defineConfig" in config_content

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!")

tests/test_js_project_root_per_function.py

@@ -0,0 +1,66 @@
+"""Test that js_project_root is recalculated per function, not cached."""
+
+from pathlib import Path
+
+from codeflash.languages.javascript.test_runner import find_node_project_root
+
+
+def test_find_node_project_root_returns_different_roots_for_different_files(tmp_path: Path) -> None:
+    """Test that find_node_project_root returns the correct root for each file."""
+    # Create main project structure
+    main_project = (tmp_path / "project").resolve()
+    main_project.mkdir()
+    (main_project / "package.json").write_text("{}", encoding="utf-8")
+    (main_project / "src").mkdir()
+    main_file = (main_project / "src" / "main.ts").resolve()
+    main_file.write_text("// main file", encoding="utf-8")
+
+    # Create extension subdirectory with its own package.json
+    extension_dir = (main_project / "extensions" / "discord").resolve()
+    extension_dir.mkdir(parents=True)
+    (extension_dir / "package.json").write_text("{}", encoding="utf-8")
+    (extension_dir / "src").mkdir()
+    extension_file = (extension_dir / "src" / "accounts.ts").resolve()
+    extension_file.write_text("// extension file", encoding="utf-8")
+
+    # Extension file should return extension directory
+    result1 = find_node_project_root(extension_file)
+    assert result1 == extension_dir, f"Expected {extension_dir}, got {result1}"
+
+    # Main file should return main project directory
+    result2 = find_node_project_root(main_file)
+    assert result2 == main_project, f"Expected {main_project}, got {result2}"
+
+    # Calling again with extension file should still return extension dir
+    result3 = find_node_project_root(extension_file)
+    assert result3 == extension_dir, f"Expected {extension_dir}, got {result3}"
+
+
+def test_js_project_root_recalculated_per_function(tmp_path: Path) -> None:
+    """Each function in a monorepo should resolve to its own nearest package.json root."""
+    # Create main project
+    main_project = (tmp_path / "project").resolve()
+    main_project.mkdir()
+    (main_project / "package.json").write_text('{"name": "main"}', encoding="utf-8")
+    (main_project / "src").mkdir()
+
+    # Create extension with its own package.json
+    extension_dir = (main_project / "extensions" / "discord").resolve()
+    extension_dir.mkdir(parents=True)
+    (extension_dir / "package.json").write_text('{"name": "discord-extension"}', encoding="utf-8")
+    (extension_dir / "src").mkdir()
+
+    extension_file = (extension_dir / "src" / "accounts.ts").resolve()
+    extension_file.write_text("export function foo() {}", encoding="utf-8")
+
+    main_file = (main_project / "src" / "commands.ts").resolve()
+    main_file.write_text("export function bar() {}", encoding="utf-8")
+
+    js_project_root_1 = find_node_project_root(extension_file)
+    assert js_project_root_1 == extension_dir
+
+    js_project_root_2 = find_node_project_root(main_file)
+    assert js_project_root_2 == main_project, (
+        f"Expected {main_project}, got {js_project_root_2}. "
+        f"Happens when js_project_root is not recalculated per function."
+    )