[FIX][JS] Avoid adding .js to relative imports in tests with jest as it's using its own resolver

codeflash/languages/javascript/instrument.py

@@ -1368,3 +1368,78 @@ def fix_mock_path(match: re.Match[str]) -> str:
         return original  # Keep original if we can't fix it
 
     return mock_pattern.sub(fix_mock_path, test_code)
+
+
+def fix_import_paths(test_code: str, test_file_path: Path, source_file_path: Path, tests_root: Path) -> str:
+    """Fix relative paths in import/require statements to be correct from the test file's location.
+
+    The AI sometimes generates import statements with paths relative to the source file
+    instead of the test file. This applies the same logic as fix_jest_mock_paths but for
+    regular imports: `import ... from '...'`, `require('...')`, and `import('...')`.
+
+    Also splits concatenated import statements onto separate lines.
+    """
+    if not test_code or not test_code.strip():
+        return test_code
+
+    import os
+
+    source_dir = source_file_path.resolve().parent
+    test_dir = test_file_path.resolve().parent
+
+    # Match: import ... from './...' or '../...',  require('./...' or '../...'),  import('./...' or '../...')
+    # Excludes jest.mock/vi.mock (handled by fix_jest_mock_paths)
+    import_pattern = re.compile(
+        r"("
+        r"(?:from\s+['\"])"  # from '...' or from "..."
+        r"|(?:require\s*\(\s*['\"])"  # require('...')
+        r"|(?:import\s*\(\s*['\"])"  # import('...')
+        r")"
+        r"(\.\./[^'\"]+|\.\/[^'\"]+)"  # relative path
+        r"(['\"])"  # closing quote
+    )
+
+    def fix_path(match: re.Match[str]) -> str:
+        original = match.group(0)
+        prefix = match.group(1)
+        rel_path = match.group(2)
+        suffix = match.group(3)
+
+        source_relative_resolved = (source_dir / rel_path).resolve()
+
+        try:
+            test_relative_resolved = (test_dir / rel_path).resolve()
+
+            if test_relative_resolved.exists() or (
+                test_relative_resolved.with_suffix(".ts").exists()
+                or test_relative_resolved.with_suffix(".js").exists()
+                or test_relative_resolved.with_suffix(".tsx").exists()
+                or test_relative_resolved.with_suffix(".jsx").exists()
+            ):
+                return original
+
+            if source_relative_resolved.exists() or (
+                source_relative_resolved.with_suffix(".ts").exists()
+                or source_relative_resolved.with_suffix(".js").exists()
+                or source_relative_resolved.with_suffix(".tsx").exists()
+                or source_relative_resolved.with_suffix(".jsx").exists()
+            ):
+                new_rel_path = Path(os.path.relpath(source_relative_resolved, test_dir)).as_posix()
+                if not new_rel_path.startswith("../") and not new_rel_path.startswith("./"):
+                    new_rel_path = f"./{new_rel_path}"
+
+                logger.debug(f"Fixed import path: {rel_path} -> {new_rel_path}")
+                return f"{prefix}{new_rel_path}{suffix}"
+
+        except (ValueError, OSError):
+            pass
+
+        return original
+
+    result = import_pattern.sub(fix_path, test_code)
+
+    # Split concatenated import/require statements onto separate lines.
+    # The AI sometimes generates: `import { a } from './x';import b from './y';`
+    # This inserts a newline before any import/require/const-require that follows a semicolon on the same line.
+    result = re.sub(r";(import\s)", r";\n\1", result)
+    return re.sub(r";(const\s+\w+\s*=\s*require\()", r";\n\1", result)

codeflash/languages/javascript/support.py

@@ -2037,6 +2037,7 @@ def process_generated_test_strings(
     ) -> tuple[str, str, str]:
         from codeflash.languages.javascript.instrument import (
             TestingMode,
+            fix_import_paths,
             fix_imports_inside_test_blocks,
             fix_jest_mock_paths,
             instrument_generated_js_test,
@@ -2058,6 +2059,11 @@ def process_generated_test_strings(
             generated_test_source, test_path, source_file, test_cfg.tests_project_rootdir
         )
 
+        # Fix relative paths in regular import/require statements
+        generated_test_source = fix_import_paths(
+            generated_test_source, test_path, source_file, test_cfg.tests_project_rootdir
+        )
+
         # Validate and fix import styles (default vs named exports)
         generated_test_source = validate_and_fix_import_style(
             generated_test_source, source_file, function_to_optimize.function_name
@@ -2070,7 +2076,8 @@ def process_generated_test_strings(
 
         # Add .js extensions to relative imports for ESM projects
         # TypeScript + ESM requires explicit .js extensions even for .ts source files
-        if project_module_system == ModuleSystem.ES_MODULE:
+        # jest uses it's own resolver so imports without the .js extension work fine
+        if project_module_system == ModuleSystem.ES_MODULE and test_cfg.test_framework != "jest":
             from codeflash.languages.javascript.module_system import add_js_extensions_to_relative_imports
 
             generated_test_source = add_js_extensions_to_relative_imports(generated_test_source)

tests/test_languages/test_javascript_instrumentation.py

@@ -865,6 +865,160 @@ def test_empty_code(self):
             assert fix_jest_mock_paths("   ", test_file, source_file, tests_dir) == "   "
 
 
+class TestFixImportPaths:
+    """Tests for fix_import_paths function."""
+
+    def test_fix_import_path_when_source_relative(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        src_dir = (tmp_path / "src" / "queue").resolve()
+        tests_dir = (tmp_path / "tests").resolve()
+        env_file = (tmp_path / "src" / "utilities" / "workerRequests.ts").resolve()
+
+        src_dir.mkdir(parents=True)
+        tests_dir.mkdir(parents=True)
+        env_file.parent.mkdir(parents=True, exist_ok=True)
+        env_file.write_text("export function sendSmtpEmail() {}", encoding="utf-8")
+
+        source_file = (src_dir / "queue.ts").resolve()
+        source_file.write_text("import { sendSmtpEmail } from '../../utilities/workerRequests';", encoding="utf-8")
+
+        test_file = (tests_dir / "test_queue.test.ts").resolve()
+
+        test_code = "import { sendSmtpEmail } from '../utilities/workerRequests';\n"
+        fixed = fix_import_paths(test_code, test_file, source_file, tests_dir)
+
+        assert "from '../src/utilities/workerRequests'" in fixed
+
+    def test_preserve_valid_import_path(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        src_dir = (tmp_path / "src").resolve()
+        tests_dir = (tmp_path / "tests").resolve()
+
+        src_dir.mkdir(parents=True)
+        tests_dir.mkdir(parents=True)
+
+        utils_file = (src_dir / "utils.ts").resolve()
+        utils_file.write_text("export const utils = {};", encoding="utf-8")
+
+        source_file = (src_dir / "main.ts").resolve()
+        source_file.write_text("", encoding="utf-8")
+        test_file = (tests_dir / "test_main.test.ts").resolve()
+
+        test_code = "import { utils } from '../src/utils';\n"
+        fixed = fix_import_paths(test_code, test_file, source_file, tests_dir)
+
+        assert "from '../src/utils'" in fixed
+
+    def test_fix_require_path(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        src_dir = (tmp_path / "src" / "queue").resolve()
+        tests_dir = (tmp_path / "tests").resolve()
+        env_file = (tmp_path / "src" / "environment.ts").resolve()
+
+        src_dir.mkdir(parents=True)
+        tests_dir.mkdir(parents=True)
+        env_file.parent.mkdir(parents=True, exist_ok=True)
+        env_file.write_text("module.exports = {};", encoding="utf-8")
+
+        source_file = (src_dir / "queue.ts").resolve()
+        source_file.write_text("", encoding="utf-8")
+        test_file = (tests_dir / "test_queue.test.ts").resolve()
+
+        test_code = "const env = require('../environment');\n"
+        fixed = fix_import_paths(test_code, test_file, source_file, tests_dir)
+
+        assert "require('../src/environment')" in fixed
+
+    def test_fix_dynamic_import_path(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        src_dir = (tmp_path / "src" / "queue").resolve()
+        tests_dir = (tmp_path / "tests").resolve()
+        env_file = (tmp_path / "src" / "environment.ts").resolve()
+
+        src_dir.mkdir(parents=True)
+        tests_dir.mkdir(parents=True)
+        env_file.parent.mkdir(parents=True, exist_ok=True)
+        env_file.write_text("export default {};", encoding="utf-8")
+
+        source_file = (src_dir / "queue.ts").resolve()
+        source_file.write_text("", encoding="utf-8")
+        test_file = (tests_dir / "test_queue.test.ts").resolve()
+
+        test_code = "const env = await import('../environment');\n"
+        fixed = fix_import_paths(test_code, test_file, source_file, tests_dir)
+
+        assert "import('../src/environment')" in fixed
+
+    def test_empty_code(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        tests_dir = (tmp_path / "tests").resolve()
+        tests_dir.mkdir()
+        source_file = (tmp_path / "src" / "main.ts").resolve()
+        test_file = (tests_dir / "test.ts").resolve()
+
+        assert fix_import_paths("", test_file, source_file, tests_dir) == ""
+        assert fix_import_paths("   ", test_file, source_file, tests_dir) == "   "
+
+    def test_does_not_touch_package_imports(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        src_dir = (tmp_path / "src").resolve()
+        tests_dir = (tmp_path / "tests").resolve()
+
+        src_dir.mkdir(parents=True)
+        tests_dir.mkdir(parents=True)
+
+        source_file = (src_dir / "main.ts").resolve()
+        source_file.write_text("", encoding="utf-8")
+        test_file = (tests_dir / "test_main.test.ts").resolve()
+
+        test_code = "import { jest } from '@jest/globals';\nimport express from 'express';\n"
+        fixed = fix_import_paths(test_code, test_file, source_file, tests_dir)
+
+        assert fixed == test_code
+
+    def test_splits_concatenated_imports(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        src_dir = (tmp_path / "src").resolve()
+        tests_dir = (tmp_path / "tests").resolve()
+
+        src_dir.mkdir(parents=True)
+        tests_dir.mkdir(parents=True)
+
+        source_file = (src_dir / "main.ts").resolve()
+        source_file.write_text("", encoding="utf-8")
+        test_file = (tests_dir / "test_main.test.ts").resolve()
+
+        test_code = "import { sendSmtpEmail } from '../src/workerRequests';import automationUtils from '../src/utils';describe('run', () => {});\n"
+        fixed = fix_import_paths(test_code, test_file, source_file, tests_dir)
+
+        assert "workerRequests';\nimport automationUtils" in fixed
+
+    def test_splits_concatenated_require(self, tmp_path: Path):
+        from codeflash.languages.javascript.instrument import fix_import_paths
+
+        src_dir = (tmp_path / "src").resolve()
+        tests_dir = (tmp_path / "tests").resolve()
+
+        src_dir.mkdir(parents=True)
+        tests_dir.mkdir(parents=True)
+
+        source_file = (src_dir / "main.ts").resolve()
+        source_file.write_text("", encoding="utf-8")
+        test_file = (tests_dir / "test_main.test.ts").resolve()
+
+        test_code = "import { foo } from '../src/foo';const bar = require('../src/bar');\n"
+        fixed = fix_import_paths(test_code, test_file, source_file, tests_dir)
+
+        assert "foo';\nconst bar = require" in fixed
+
+
 class TestFunctionCallsInStrings:
     """Tests for skipping function calls inside string literals."""