fix: support Mocha CJS projects and sanitize incorrect framework imports

Three related fixes for Mocha test generation in CommonJS projects:

1. inject_test_globals() now accepts module_system param — emits
   `require('node:assert/strict')` for CJS instead of ESM import syntax
2. ensure_module_system_compatibility() now converts ESM→CJS even when
   the source has mixed imports (was skipping when both ESM and CJS were
   detected, leaving the ESM import from inject_test_globals unconverted)
3. New sanitize_mocha_imports() strips vitest/jest/@jest/globals imports
   that the AI sometimes generates for Mocha projects — Mocha provides
   describe/it/before*/after* as globals

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

codeflash/languages/javascript/edit_tests.py

@@ -214,24 +214,30 @@ def normalize_codeflash_imports(source: str) -> str:
 
 
 # Author: ali <mohammed18200118@gmail.com>
-def inject_test_globals(generated_tests: GeneratedTestsList, test_framework: str = "jest") -> GeneratedTestsList:
+def inject_test_globals(
+    generated_tests: GeneratedTestsList, test_framework: str = "jest", module_system: str = "esm"
+) -> GeneratedTestsList:
     # TODO: inside the prompt tell the llm if it should import jest functions or it's already injected in the global window
     """Inject test globals into all generated tests.
 
     Args:
         generated_tests: List of generated tests.
         test_framework: The test framework being used ("jest", "vitest", or "mocha").
+        module_system: The module system ("esm" or "commonjs").
 
     Returns:
         Generated tests with test globals injected.
 
     """
-    # we only inject test globals for esm modules
+    is_cjs = module_system == "commonjs"
     # Use vitest imports for vitest projects, jest imports for jest projects
     if test_framework == "vitest":
         global_import = "import { vi, describe, it, expect, beforeEach, afterEach, beforeAll, test } from 'vitest'\n"
     elif test_framework == "mocha":
-        global_import = "import assert from 'node:assert/strict';\n"
+        if is_cjs:
+            global_import = "const assert = require('node:assert/strict');\n"
+        else:
+            global_import = "import assert from 'node:assert/strict';\n"
     else:
         # Default to jest imports for jest and other frameworks
         global_import = (
@@ -245,6 +251,35 @@ def inject_test_globals(generated_tests: GeneratedTestsList, test_framework: str
     return generated_tests
 
 
+_VITEST_IMPORT_RE = re.compile(r"^.*import\s+\{[^}]*\}\s+from\s+['\"]vitest['\"].*\n?", re.MULTILINE)
+_JEST_GLOBALS_IMPORT_RE = re.compile(r"^.*import\s+\{[^}]*\}\s+from\s+['\"]@jest/globals['\"].*\n?", re.MULTILINE)
+_MOCHA_REQUIRE_RE = re.compile(
+    r"^.*(?:const|let|var)\s+\{[^}]*\}\s*=\s*require\s*\(\s*['\"]mocha['\"]\s*\).*\n?", re.MULTILINE
+)
+_VITEST_COMMENT_RE = re.compile(r"^.*//.*vitest imports.*\n?", re.MULTILINE | re.IGNORECASE)
+
+
+def sanitize_mocha_imports(source: str) -> str:
+    """Remove vitest/jest/mocha-require imports from Mocha test source.
+
+    The AI service sometimes generates vitest or jest-style imports when the
+    framework is mocha. Mocha provides describe/it/before*/after* as globals,
+    so these imports must be removed. Also removes ``require('mocha')``
+    destructures since Mocha doesn't export those.
+
+    Args:
+        source: Generated test source code.
+
+    Returns:
+        Source with incorrect framework imports stripped.
+
+    """
+    source = _VITEST_IMPORT_RE.sub("", source)
+    source = _JEST_GLOBALS_IMPORT_RE.sub("", source)
+    source = _MOCHA_REQUIRE_RE.sub("", source)
+    return _VITEST_COMMENT_RE.sub("", source)
+
+
 # Author: ali <mohammed18200118@gmail.com>
 def disable_ts_check(generated_tests: GeneratedTestsList) -> GeneratedTestsList:
     """Disable TypeScript type checking in all generated tests.

codeflash/languages/javascript/module_system.py

@@ -439,8 +439,8 @@ def ensure_module_system_compatibility(code: str, target_module_system: str, pro
         logger.debug("Converting CommonJS require statements to ES Module syntax")
         return convert_commonjs_to_esm(code)
 
-    if target_module_system == ModuleSystem.COMMONJS and is_esm and not is_commonjs:
-        logger.debug("Converting ES Module to CommonJS syntax")
+    if target_module_system == ModuleSystem.COMMONJS and is_esm:
+        logger.debug("Converting ES Module imports to CommonJS syntax")
         return convert_esm_to_commonjs(code)
 
     logger.debug("No module system conversion needed")

codeflash/languages/javascript/support.py

@@ -1787,12 +1787,32 @@ class JavaScriptSupport:
             disable_ts_check,
             inject_test_globals,
             normalize_generated_tests_imports,
+            sanitize_mocha_imports,
         )
         from codeflash.languages.javascript.module_system import detect_module_system
+        from codeflash.models.models import GeneratedTests as GeneratedTestsModel
+        from codeflash.models.models import GeneratedTestsList
+
+        # For Mocha, strip vitest/jest/require('mocha') imports the AI may have generated
+        if test_framework == "mocha":
+            sanitized = []
+            for test in generated_tests.generated_tests:
+                sanitized.append(
+                    GeneratedTestsModel(
+                        generated_original_test_source=sanitize_mocha_imports(test.generated_original_test_source),
+                        instrumented_behavior_test_source=sanitize_mocha_imports(
+                            test.instrumented_behavior_test_source
+                        ),
+                        instrumented_perf_test_source=sanitize_mocha_imports(test.instrumented_perf_test_source),
+                        behavior_file_path=test.behavior_file_path,
+                        perf_file_path=test.perf_file_path,
+                    )
+                )
+            generated_tests = GeneratedTestsList(generated_tests=sanitized)
 
         module_system = detect_module_system(project_root, source_file_path)
-        if module_system == "esm":
-            generated_tests = inject_test_globals(generated_tests, test_framework)
+        if module_system == "esm" or test_framework == "mocha":
+            generated_tests = inject_test_globals(generated_tests, test_framework, module_system)
         if self.language == Language.TYPESCRIPT:
             generated_tests = disable_ts_check(generated_tests)
         return normalize_generated_tests_imports(generated_tests)

tests/test_languages/test_mocha_runner.py

@@ -456,6 +456,142 @@ class TestRunMochaBenchmarkingTests:
             assert env.get("CODEFLASH_PERF_STABILITY_CHECK") == "false"
 
 
+class TestSanitizeMochaImports:
+    """Tests for stripping wrong framework imports from Mocha tests."""
+
+    def test_strips_vitest_import(self):
+        from codeflash.languages.javascript.edit_tests import sanitize_mocha_imports
+
+        source = "import { describe, test, expect, vi } from 'vitest'\nconst x = 1;\n"
+        result = sanitize_mocha_imports(source)
+        assert "vitest" not in result
+        assert "const x = 1;" in result
+
+    def test_strips_jest_globals_import(self):
+        from codeflash.languages.javascript.edit_tests import sanitize_mocha_imports
+
+        source = "import { jest, describe, it, expect } from '@jest/globals'\nconst x = 1;\n"
+        result = sanitize_mocha_imports(source)
+        assert "@jest/globals" not in result
+        assert "const x = 1;" in result
+
+    def test_strips_mocha_require(self):
+        from codeflash.languages.javascript.edit_tests import sanitize_mocha_imports
+
+        source = "const { describe, it, expect } = require('mocha');\nconst x = 1;\n"
+        result = sanitize_mocha_imports(source)
+        assert "require('mocha')" not in result
+        assert "const x = 1;" in result
+
+    def test_strips_vitest_comment(self):
+        from codeflash.languages.javascript.edit_tests import sanitize_mocha_imports
+
+        source = "// vitest imports (REQUIRED for vitest)\nimport { describe } from 'vitest'\nconst x = 1;\n"
+        result = sanitize_mocha_imports(source)
+        assert "vitest" not in result
+        assert "const x = 1;" in result
+
+    def test_preserves_unrelated_imports(self):
+        from codeflash.languages.javascript.edit_tests import sanitize_mocha_imports
+
+        source = "const sinon = require('sinon');\nconst assert = require('node:assert/strict');\n"
+        result = sanitize_mocha_imports(source)
+        assert "sinon" in result
+        assert "node:assert/strict" in result
+
+
+class TestInjectTestGlobalsModuleSystem:
+    """Tests for inject_test_globals with different module systems."""
+
+    def test_mocha_esm_uses_import(self):
+        from codeflash.languages.javascript.edit_tests import inject_test_globals
+        from codeflash.models.models import GeneratedTests, GeneratedTestsList
+
+        tests = GeneratedTestsList(
+            generated_tests=[
+                GeneratedTests(
+                    generated_original_test_source="describe('test', () => {});",
+                    instrumented_behavior_test_source="describe('test', () => {});",
+                    instrumented_perf_test_source="describe('test', () => {});",
+                    behavior_file_path=Path("test.test.js"),
+                    perf_file_path=Path("test.perf.test.js"),
+                )
+            ]
+        )
+
+        result = inject_test_globals(tests, test_framework="mocha", module_system="esm")
+        assert "import assert from 'node:assert/strict'" in result.generated_tests[0].generated_original_test_source
+
+    def test_mocha_cjs_uses_require(self):
+        from codeflash.languages.javascript.edit_tests import inject_test_globals
+        from codeflash.models.models import GeneratedTests, GeneratedTestsList
+
+        tests = GeneratedTestsList(
+            generated_tests=[
+                GeneratedTests(
+                    generated_original_test_source="describe('test', () => {});",
+                    instrumented_behavior_test_source="describe('test', () => {});",
+                    instrumented_perf_test_source="describe('test', () => {});",
+                    behavior_file_path=Path("test.test.js"),
+                    perf_file_path=Path("test.perf.test.js"),
+                )
+            ]
+        )
+
+        result = inject_test_globals(tests, test_framework="mocha", module_system="commonjs")
+        src = result.generated_tests[0].generated_original_test_source
+        assert "const assert = require('node:assert/strict')" in src
+        assert "import assert" not in src
+
+    def test_vitest_always_uses_import(self):
+        from codeflash.languages.javascript.edit_tests import inject_test_globals
+        from codeflash.models.models import GeneratedTests, GeneratedTestsList
+
+        tests = GeneratedTestsList(
+            generated_tests=[
+                GeneratedTests(
+                    generated_original_test_source="describe('test', () => {});",
+                    instrumented_behavior_test_source="describe('test', () => {});",
+                    instrumented_perf_test_source="describe('test', () => {});",
+                    behavior_file_path=Path("test.test.js"),
+                    perf_file_path=Path("test.perf.test.js"),
+                )
+            ]
+        )
+
+        result = inject_test_globals(tests, test_framework="vitest", module_system="commonjs")
+        assert "from 'vitest'" in result.generated_tests[0].generated_original_test_source
+
+
+class TestEnsureModuleSystemCompatibilityMixed:
+    """Tests for ensure_module_system_compatibility with mixed ESM+CJS code."""
+
+    def test_converts_imports_in_mixed_code_to_cjs(self):
+        from codeflash.languages.javascript.module_system import ensure_module_system_compatibility
+
+        # Code with both import (from inject_test_globals) and require (from backend)
+        code = "import assert from 'node:assert/strict';\nconst { foo } = require('./module');\n"
+        result = ensure_module_system_compatibility(code, "commonjs")
+        assert "require('node:assert/strict')" in result
+        assert "import assert" not in result
+
+    def test_converts_require_in_mixed_code_to_esm(self):
+        from codeflash.languages.javascript.module_system import ensure_module_system_compatibility
+
+        code = "import { describe } from 'vitest';\nconst foo = require('./module');\n"
+        result = ensure_module_system_compatibility(code, "esm")
+        assert "require" not in result
+        assert "import" in result
+
+    def test_pure_esm_to_cjs(self):
+        from codeflash.languages.javascript.module_system import ensure_module_system_compatibility
+
+        code = "import assert from 'node:assert/strict';\nimport { foo } from './module';\n"
+        result = ensure_module_system_compatibility(code, "commonjs")
+        assert "require('node:assert/strict')" in result
+        assert "import" not in result
+
+
 class TestRunMochaLineProfileTests:
     """Tests for running Mocha line profile tests with mocked subprocess."""