Skip to content

alias support and vitest imports #1266

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Saga4 merged 2 commits into from
Feb 3, 2026
Merged

alias support and vitest imports #1266

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
83113d1
alias support and vitest imports
Saga4 Feb 2, 2026
824b2ec
Merge branch 'main' into alias_support_in_ts
Saga4 Feb 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions codeflash/cli_cmds/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,8 +288,9 @@ def process_pyproject_config(args: Namespace) -> Namespace:
normalized_ignore_paths = []
for path in args.ignore_paths:
path_obj = Path(path)
assert path_obj.exists(), f"ignore-paths config must be a valid path. Path {path} does not exist"
normalized_ignore_paths.append(path_obj.resolve())
if path_obj.exists():
normalized_ignore_paths.append(path_obj.resolve())
# Silently skip non-existent paths (e.g., .next, dist before build)
args.ignore_paths = normalized_ignore_paths
# Project root path is one level above the specified directory, because that's where the module can be imported from
args.module_root = Path(args.module_root).resolve()
Expand Down
52 changes: 51 additions & 1 deletion codeflash/languages/javascript/import_resolver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,15 @@ def _resolve_module_path(self, module_path: str, source_dir: Path) -> Path | Non
if module_path.startswith("/"):
return self._resolve_absolute_import(module_path)

# Handle @/ path alias (common in Next.js/TypeScript projects)
# @/ maps to the project root
if module_path.startswith("@/"):
return self._resolve_path_alias(module_path[2:]) # Strip @/

# Handle ~/ path alias (another common pattern)
if module_path.startswith("~/"):
return self._resolve_path_alias(module_path[2:]) # Strip ~/

# Bare imports (e.g., 'lodash') are external packages
return None

Expand Down Expand Up @@ -197,6 +206,38 @@ def _resolve_absolute_import(self, module_path: str) -> Path | None:

return None

def _resolve_path_alias(self, module_path: str) -> Path | None:
"""Resolve path alias imports like @/utils or ~/lib/helper.

Args:
module_path: The import path without the alias prefix.

Returns:
Resolved absolute path, or None if not found.

"""
# Treat as relative to project root
base_path = (self.project_root / module_path).resolve()

# Check if path is within project
try:
base_path.relative_to(self.project_root)
except ValueError:
logger.debug("Path alias resolves outside project root: %s", base_path)
return None

# Try adding extensions
resolved = self._try_extensions(base_path)
if resolved:
return resolved

# Try as directory with index file
resolved = self._try_index_file(base_path)
if resolved:
return resolved

return None

def _try_extensions(self, base_path: Path) -> Path | None:
"""Try adding various extensions to find the actual file.

Expand Down Expand Up @@ -267,10 +308,19 @@ def _is_external_package(self, module_path: str) -> bool:
if module_path.startswith("/"):
return False

# @/ is a common path alias in Next.js/TypeScript projects mapping to project root
# These are internal imports, not external npm packages
if module_path.startswith("@/"):
return False

# ~/ is another common path alias pattern
if module_path.startswith("~/"):
return False

# Bare imports without ./ or ../ are external packages
# This includes:
# - 'lodash'
# - '@company/utils'
# - '@company/utils' (scoped npm packages)
# - 'react'
# - 'fs' (Node.js built-ins)
return True
Expand Down
67 changes: 67 additions & 0 deletions codeflash/languages/javascript/module_system.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,3 +330,70 @@ def ensure_module_system_compatibility(code: str, target_module_system: str) ->
return convert_esm_to_commonjs(code)

return code


def ensure_vitest_imports(code: str, test_framework: str) -> str:
"""Ensure vitest test globals are imported when using vitest framework.

Vitest by default does not enable globals (describe, test, expect, etc.),
so they must be explicitly imported. This function adds the import if missing.

Args:
code: JavaScript/TypeScript test code.
test_framework: The test framework being used (vitest, jest, mocha).

Returns:
Code with vitest imports added if needed.

"""
if test_framework != "vitest":
return code

# Check if vitest imports already exist
if "from 'vitest'" in code or 'from "vitest"' in code:
return code

# Check if the code uses test functions that need to be imported
test_globals = ["describe", "test", "it", "expect", "vi", "beforeEach", "afterEach", "beforeAll", "afterAll"]
needs_import = any(f"{global_name}(" in code or f"{global_name} (" in code for global_name in test_globals)

if not needs_import:
return code

# Determine which globals are actually used in the code
used_globals = [g for g in test_globals if f"{g}(" in code or f"{g} (" in code]
if not used_globals:
return code

# Build the import statement
import_statement = f"import {{ {', '.join(used_globals)} }} from 'vitest';\n"

# Find the first line that isn't a comment or empty
lines = code.split("\n")
insert_index = 0
for i, line in enumerate(lines):
stripped = line.strip()
if stripped and not stripped.startswith("//") and not stripped.startswith("/*") and not stripped.startswith("*"):
# Check if this line is an import/require - insert after imports
if stripped.startswith("import ") or stripped.startswith("const ") or stripped.startswith("let "):
continue
insert_index = i
break
insert_index = i + 1

# Find the last import line to insert after it
last_import_index = -1
for i, line in enumerate(lines):
stripped = line.strip()
if stripped.startswith("import ") and "from " in stripped:
last_import_index = i

Comment on lines +358 to +390
Copy link
Contributor

@codeflash-ai codeflash-ai bot Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 23% (0.23x) speedup for ensure_vitest_imports in codeflash/languages/javascript/module_system.py

⏱️ Runtime : 828 microseconds → 673 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 23% runtime improvement by eliminating redundant scanning of the code string and consolidating loop iterations over the code lines.

Key optimizations:

  1. Single-pass global detection: The original code scanned for test globals twice—once with any() to check if imports are needed, then again with a list comprehension to collect used globals. The optimized version computes used_globals in a single pass and derives needs_import from whether that list is empty using bool(used_globals). This eliminates one complete scan through the code string for all 9 test globals.

  2. Merged loop for insertion logic: The original code had two separate loops over lines:

    • First loop: Find insert_index (where to insert if no imports exist)
    • Second loop: Find last_import_index (the last existing import line)

    The optimized version merges these into a single loop, checking for both conditions simultaneously. This cuts the number of line iterations in half, reducing overhead from enumerate(), strip() calls, and string comparisons.

Performance characteristics:

The test results show the optimization is most effective for:

  • Large files (27-41% faster): The benefit compounds with file size since we avoid redundant string scans and line iterations
  • Multiple globals usage (30-48% faster): Cases using test, expect, it, afterAll, etc. benefit significantly from the single-pass detection
  • Files with many imports (10-18% faster): Merged loop reduces overhead when scanning through existing import statements

The optimization maintains identical behavior—all test cases pass with the same output while running faster across all scenarios, from simple cases (~6-35% faster) to complex nested structures (~41% faster).

✅ Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests ✅ 49 Passed
⏪ Replay Tests 🔘 None Found
🔎 Concolic Coverage Tests 🔘 None Found
📊 Tests Coverage 96.9%
🌀 Click to see Generated Regression Tests 
import logging

# imports
import pytest  # used for our unit tests
from codeflash.languages.javascript.module_system import ensure_vitest_imports

def test_non_vitest_framework_returns_same():
    # If the framework is not 'vitest', the code must be returned unchanged.
    original = "describe('a', () => {});"
    codeflash_output = ensure_vitest_imports(original, test_framework="jest"); result = codeflash_output # 861ns -> 842ns (2.26% faster)

def test_no_globals_no_insertion():
    # When the code doesn't use any vitest globals, nothing should be added.
    original = "// file with comments only\nconst x = 1;\nfunction foo() { return x; }\n"
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 4.98μs -> 4.59μs (8.50% faster)

def test_already_imported_single_quotes():
    # If the code already imports from 'vitest' (single quotes), return unchanged.
    original = "import { describe } from 'vitest';\ndescribe('a', () => {});"
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 1.04μs -> 981ns (6.22% faster)

def test_already_imported_double_quotes():
    # If the code already imports from "vitest" (double quotes), return unchanged.
    original = 'import { test } from "vitest";\ntest("works", () => {});'
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 1.06μs -> 1.05μs (0.951% faster)

def test_basic_add_import_at_top_for_describe_and_test():
    # When using vitest and code contains globals, the function should add the import.
    original = "describe('suite', () => {\n  test('case', () => {});\n});"
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 10.4μs -> 8.45μs (22.9% faster)

def test_insert_after_last_import_line():
    # If the code has existing import ... from ... lines, the vitest import
    # should be inserted after the last such import.
    original = (
        "import foo from 'foo';\n"
        "import bar from 'bar';\n"
        "// some comment\n"
        "describe('a', () => {});\n"
    )
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 11.7μs -> 9.95μs (17.3% faster)
    lines = result.splitlines()

def test_insert_after_const_or_let_lines():
    # Lines beginning with const/let should be skipped when determining insertion
    # point, so the vitest import should be inserted after them.
    original = "const x = require('x');\ndescribe('a', () => {});"
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 9.92μs -> 8.32μs (19.1% faster)
    lines = result.splitlines()

def test_detects_usage_in_string_literals_and_inserts():
    # The function currently checks substrings, so if a global name appears inside
    # a string literal like "describe(", it will still trigger an import.
    original = "const s = \"this mentions describe(\";\nconsole.log(s);"
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 9.91μs -> 8.28μs (19.7% faster)

def test_handles_space_before_parenthesis():
    # The detection should handle both 'it(' and 'it (' (with a space).
    original = "it ('should pass', () => {});"
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 9.73μs -> 7.39μs (31.6% faster)

def test_used_globals_order_is_preserved():
    # When multiple globals are used, their order in the import should follow
    # the original test_globals list defined in the function.
    # Here we use 'vi' and 'expect' so import should be "expect, vi"
    original = "vi(functionStub);\nexpect(1).toBe(1);\n"
    # Note: order in the code is vi first, expect second, but used_globals must follow
    # the fixed ordering inside the function (expect before vi).
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 10.4μs -> 7.66μs (35.3% faster)
    # Find the import line
    import_lines = [l for l in result.splitlines() if l.startswith("import {")]
    import_line = import_lines[0]

def test_large_scale_insertion_efficiency_and_correctness():
    # Large-scale-ish test: create a file with several hundred lines (under 1000)
    # and a few scattered usages of vitest globals. Ensure the import is added
    # exactly once and lists only the used globals (in the canonical order).
    fragments = []
    # Add 250 noop lines
    for i in range(250):
        fragments.append(f"// filler line {i}")
    # Add some import statements to ensure insertion after last import
    fragments.append("import alpha from 'alpha';")
    fragments.append("import beta from 'beta';")
    # Add more filler
    for i in range(250, 500):
        fragments.append(f"const val{i} = {i};")
    # Use a few vitest globals scattered through the file
    fragments.append("beforeEach(() => {});")
    fragments.append("someOtherLine();")
    fragments.append("afterAll(() => {});")
    fragments.append("console.log('done');")
    # Join into a big code string
    big_code = "\n".join(fragments)
    codeflash_output = ensure_vitest_imports(big_code, test_framework="vitest"); result = codeflash_output # 346μs -> 271μs (27.4% faster)
    # Ensure the import was inserted after the last explicit import ('import beta from ...')
    lines = result.splitlines()
    # Find index of "import beta from 'beta';"
    beta_idx = next(i for i, l in enumerate(lines) if l.strip() == "import beta from 'beta';")

def test_no_duplicate_insertion_on_multiple_global_occurrences():
    # If globals appear multiple times, the import should still list each used global once.
    original = (
        "describe('one', () => {});\n"
        "describe('two', () => {});\n"
        "test('a', () => {});\n"
        "test('b', () => {});\n"
    )
    codeflash_output = ensure_vitest_imports(original, test_framework="vitest"); result = codeflash_output # 10.9μs -> 8.73μs (25.0% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
import pytest
from codeflash.languages.javascript.module_system import ensure_vitest_imports

def test_non_vitest_framework_returns_code_unchanged():
    """Test that code is returned unchanged when test_framework is not 'vitest'."""
    code = "describe('test', () => { test('works', () => { expect(1).toBe(1); }); });"
    codeflash_output = ensure_vitest_imports(code, "jest"); result = codeflash_output # 511ns -> 521ns (1.92% slower)

def test_jest_framework_ignores_vitest_needs():
    """Test that jest framework is not modified even if vitest imports are missing."""
    code = "describe('suite', () => { test('case', () => {}); });"
    codeflash_output = ensure_vitest_imports(code, "jest"); result = codeflash_output # 470ns -> 441ns (6.58% faster)

def test_mocha_framework_returns_unchanged():
    """Test that mocha framework code is returned unchanged."""
    code = "describe('mocha test', () => { it('works', () => {}); });"
    codeflash_output = ensure_vitest_imports(code, "mocha"); result = codeflash_output # 441ns -> 421ns (4.75% faster)

def test_vitest_with_existing_import_single_quotes():
    """Test that existing vitest imports with single quotes are recognized."""
    code = "import { describe, test } from 'vitest';\ndescribe('test', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 731ns -> 771ns (5.19% slower)

def test_vitest_with_existing_import_double_quotes():
    """Test that existing vitest imports with double quotes are recognized."""
    code = 'import { describe, test } from "vitest";\ndescribe("test", () => {});'
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 822ns -> 872ns (5.73% slower)

def test_vitest_adds_import_for_describe():
    """Test that vitest import is added when describe is used."""
    code = "describe('test', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.30μs -> 7.49μs (24.1% faster)

def test_vitest_adds_import_for_test():
    """Test that vitest import is added when test is used."""
    code = "test('my test', () => { expect(true).toBe(true); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.39μs -> 7.17μs (30.9% faster)

def test_vitest_adds_import_for_expect():
    """Test that vitest import is added when expect is used."""
    code = "expect(5).toBe(5);"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.23μs -> 6.84μs (34.8% faster)

def test_vitest_adds_import_for_it():
    """Test that vitest import is added when it is used."""
    code = "it('should work', () => { expect(true).toBe(true); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.48μs -> 7.30μs (29.8% faster)

def test_vitest_no_import_if_no_globals_used():
    """Test that no import is added if test globals are not used."""
    code = "const x = 5; console.log(x);"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.82μs -> 3.56μs (7.31% faster)

def test_vitest_handles_multiple_globals():
    """Test that multiple vitest globals are included in a single import."""
    code = "describe('suite', () => { test('case', () => { expect(1).toBe(1); }); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.26μs -> 7.61μs (21.6% faster)

def test_vitest_with_space_before_parenthesis():
    """Test that globals with space before parenthesis are recognized."""
    code = "describe ('test', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 8.74μs -> 7.03μs (24.2% faster)

def test_vitest_with_multiple_spaces_before_parenthesis():
    """Test that globals with multiple spaces before parenthesis are recognized."""
    code = "test  ('my test', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.90μs -> 3.66μs (6.56% faster)

def test_vitest_ignores_globals_in_comments():
    """Test that globals mentioned in comments don't trigger imports."""
    code = "// describe('test', () => {});\nconst x = 5;"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.40μs -> 7.83μs (20.0% faster)

def test_vitest_ignores_globals_in_block_comments():
    """Test that globals in block comments are ignored."""
    code = "/* describe('test', () => {}); */\nconst x = 5;"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 8.97μs -> 7.41μs (20.9% faster)

def test_vitest_with_leading_comments():
    """Test that import is placed correctly after leading comments."""
    code = "// This is a test file\n// Another comment\ndescribe('test', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.88μs -> 8.27μs (19.5% faster)
    lines = result.split("\n")

def test_vitest_with_existing_imports_before_code():
    """Test that vitest import is inserted after existing imports."""
    code = "import axios from 'axios';\ndescribe('test', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.74μs -> 8.20μs (18.8% faster)
    lines = result.split("\n")
    # Vitest import should come after
    vitest_line_index = next(i for i, line in enumerate(lines) if "vitest" in line)

def test_vitest_with_beforeeach():
    """Test that beforeEach is recognized and imported."""
    code = "beforeEach(() => { setup(); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.77μs -> 6.91μs (41.3% faster)

def test_vitest_with_aftereach():
    """Test that afterEach is recognized and imported."""
    code = "afterEach(() => { cleanup(); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.70μs -> 7.01μs (38.3% faster)

def test_vitest_with_beforeall():
    """Test that beforeAll is recognized and imported."""
    code = "beforeAll(() => { init(); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 8.84μs -> 6.97μs (26.7% faster)

def test_vitest_with_afterall():
    """Test that afterAll is recognized and imported."""
    code = "afterAll(() => { destroy(); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.91μs -> 6.69μs (48.1% faster)

def test_vitest_with_vi_mock():
    """Test that vi mock utility is recognized and imported."""
    code = "vi.mock('./module');"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.90μs -> 3.52μs (10.8% faster)

def test_vitest_empty_code():
    """Test handling of empty code string."""
    code = ""
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.01μs -> 2.69μs (12.0% faster)

def test_vitest_only_whitespace():
    """Test handling of code with only whitespace."""
    code = "   \n  \n   "
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.48μs -> 3.17μs (9.79% faster)

def test_vitest_global_not_as_function_call():
    """Test that globals not followed by parenthesis are not imported."""
    code = "const test = 5; const describe = 'string';"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 3.90μs -> 3.63μs (7.44% faster)

def test_vitest_case_sensitive_globals():
    """Test that global matching is case-sensitive."""
    code = "Describe('test', () => {}); Test('case', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 4.27μs -> 4.08μs (4.66% faster)

def test_vitest_mixed_existing_and_needed_imports():
    """Test handling when some imports exist and some are needed."""
    code = "import { describe } from 'vitest';\ntest('case', () => {});"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 701ns -> 731ns (4.10% slower)

def test_vitest_with_line_continuation():
    """Test code with line continuations and multiple statements."""
    code = "describe(\n  'test',\n  () => {\n    test('case', () => {});\n  }\n);"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 10.9μs -> 8.40μs (30.1% faster)

def test_vitest_preserves_original_content():
    """Test that the original code content is preserved after adding imports."""
    code = "describe('test', () => { expect(true).toBe(true); });"
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 9.03μs -> 7.48μs (20.6% faster)

def test_vitest_large_code_file_with_many_tests():
    """Test handling of a large file with many test cases."""
    # Generate a code string with 50 test cases
    code_parts = []
    for i in range(50):
        code_parts.append(f"test('test case {i}', () => {{ expect({i}).toBe({i}); }});")
    code = "\n".join(code_parts)
    
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 40.9μs -> 31.9μs (28.2% faster)

def test_vitest_large_file_with_multiple_import_types():
    """Test handling of a large file with many existing imports."""
    # Create code with many imports before test code
    imports = []
    for i in range(50):
        imports.append(f"import module{i} from 'module{i}';")
    imports.append("describe('suite', () => { test('case', () => {}); });")
    code = "\n".join(imports)
    
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 45.9μs -> 41.7μs (10.1% faster)
    # All original imports should be preserved
    for i in range(50):
        pass

def test_vitest_large_code_with_mixed_globals():
    """Test handling of code that uses all supported vitest globals."""
    globals_used = ["describe", "test", "it", "expect", "vi", "beforeEach", "afterEach", "beforeAll", "afterAll"]
    code_parts = []
    for global_name in globals_used:
        code_parts.append(f"{global_name}(() => {{}});")
    code = "\n".join(code_parts)
    
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 10.1μs -> 7.61μs (32.1% faster)
    # All globals should be in the import
    for global_name in globals_used:
        pass

def test_vitest_large_file_no_modifications_needed():
    """Test performance with large file that doesn't need modifications."""
    # Create a large file with no vitest globals
    code_parts = []
    for i in range(100):
        code_parts.append(f"const var{i} = {i};")
    code = "\n".join(code_parts)
    
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 19.4μs -> 19.1μs (1.79% faster)

def test_vitest_complex_nested_code_structure():
    """Test handling of deeply nested code structures with multiple globals."""
    code = """
    describe('outer suite', () => {
        beforeAll(() => {});
        afterAll(() => {});
        
        describe('inner suite', () => {
            beforeEach(() => {});
            afterEach(() => {});
            
            test('test 1', () => {
                expect(1).toBe(1);
                vi.mock('./mod');
            });
            
            it('test 2', () => {
                expect(2).toBe(2);
            });
        });
    });
    """
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 15.7μs -> 11.1μs (41.0% faster)

def test_vitest_code_with_long_lines():
    """Test handling of code with very long lines."""
    long_line = "describe('test', () => { " + "test('case', () => { expect(x).toBe(y); }); " * 50 + "});"
    code = long_line
    
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 22.4μs -> 20.7μs (8.08% faster)

def test_vitest_performance_with_many_comment_lines():
    """Test that performance is acceptable with many comment lines."""
    comments = ["// This is a comment"] * 100
    code_lines = comments + ["describe('test', () => { test('case', () => {}); });"]
    code = "\n".join(code_lines)
    
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 52.8μs -> 46.6μs (13.3% faster)

def test_vitest_repeated_globals_not_duplicated():
    """Test that repeated use of same global doesn't create duplicates in import."""
    code = """
    describe('suite', () => {});
    describe('another suite', () => {});
    test('case1', () => {});
    test('case2', () => {});
    test('case3', () => {});
    """
    codeflash_output = ensure_vitest_imports(code, "vitest"); result = codeflash_output # 11.9μs -> 9.34μs (27.6% faster)
    # Count occurrences of describe and test in import
    lines = result.split("\n")
    import_line = next(line for line in lines if "import {" in line and "vitest" in line)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To test or edit this optimization locally git merge codeflash/optimize-pr1266-2026-02-03T06.54.35

Click to see suggested changes
Suggested change
needs_import = any(f"{global_name}(" in code or f"{global_name} (" in code for global_name in test_globals)
if not needs_import:
return code
# Determine which globals are actually used in the code
used_globals = [g for g in test_globals if f"{g}(" in code or f"{g} (" in code]
if not used_globals:
return code
# Build the import statement
import_statement = f"import {{ {', '.join(used_globals)} }} from 'vitest';\n"
# Find the first line that isn't a comment or empty
lines = code.split("\n")
insert_index = 0
for i, line in enumerate(lines):
stripped = line.strip()
if stripped and not stripped.startswith("//") and not stripped.startswith("/*") and not stripped.startswith("*"):
# Check if this line is an import/require - insert after imports
if stripped.startswith("import ") or stripped.startswith("const ") or stripped.startswith("let "):
continue
insert_index = i
break
insert_index = i + 1
# Find the last import line to insert after it
last_import_index = -1
for i, line in enumerate(lines):
stripped = line.strip()
if stripped.startswith("import ") and "from " in stripped:
last_import_index = i
# Compute used globals in a single pass to avoid scanning code twice
used_globals = [g for g in test_globals if f"{g}(" in code or f"{g} (" in code]
needs_import = bool(used_globals)
if not needs_import:
return code
# Build the import statement
import_statement = f"import {{ {', '.join(used_globals)} }} from 'vitest';\n"
# Find the first line that isn't a comment or empty
lines = code.split("\n")
insert_index = 0
last_import_index = -1
for i, line in enumerate(lines):
stripped = line.strip()
if stripped.startswith("import ") and "from " in stripped:
last_import_index = i
if stripped and not stripped.startswith("//") and not stripped.startswith("/*") and not stripped.startswith("*"):
# Check if this line is an import/require - insert after imports
if stripped.startswith("import ") or stripped.startswith("const ") or stripped.startswith("let "):
continue
insert_index = i
break
insert_index = i + 1

Static Badge

Copy link
Contributor Author

@Saga4 Saga4 Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Too large diff

if last_import_index >= 0:
# Insert after the last import
lines.insert(last_import_index + 1, import_statement.rstrip())
else:
# Insert at the beginning (after any leading comments)
lines.insert(insert_index, import_statement.rstrip())

logger.debug("Added vitest imports: %s", used_globals)
return "\n".join(lines)
4 changes: 3 additions & 1 deletion codeflash/languages/javascript/support.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,9 @@ def default_file_extension(self) -> str:
@property
def test_framework(self) -> str:
"""Primary test framework for JavaScript."""
return "jest"
from codeflash.languages.test_framework import get_js_test_framework_or_default

return get_js_test_framework_or_default()

@property
def comment_prefix(self) -> str:
Expand Down
7 changes: 5 additions & 2 deletions codeflash/verification/verification_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,10 +86,13 @@ class TestConfig:
def test_framework(self) -> str:
"""Returns the appropriate test framework based on language.

Returns 'jest' for JavaScript/TypeScript, 'pytest' for Python (default).
For JavaScript/TypeScript: uses the configured framework (vitest, jest, or mocha).
For Python: uses pytest as default.
"""
if is_javascript():
return "jest"
from codeflash.languages.test_framework import get_js_test_framework_or_default

return get_js_test_framework_or_default()
return "pytest"

def set_language(self, language: str) -> None:
Expand Down
8 changes: 7 additions & 1 deletion codeflash/verification/verifier.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,10 @@ def generate_tests(
instrument_generated_js_test,
validate_and_fix_import_style,
)
from codeflash.languages.javascript.module_system import ensure_module_system_compatibility
from codeflash.languages.javascript.module_system import (
ensure_module_system_compatibility,
ensure_vitest_imports,
)

source_file = Path(function_to_optimize.file_path)

Expand All @@ -81,6 +84,9 @@ def generate_tests(
# Convert module system if needed (e.g., CommonJS -> ESM for ESM projects)
generated_test_source = ensure_module_system_compatibility(generated_test_source, project_module_system)

# Ensure vitest imports are present when using vitest framework
generated_test_source = ensure_vitest_imports(generated_test_source, test_cfg.test_framework)

# Instrument for behavior verification (writes to SQLite)
instrumented_behavior_test_source = instrument_generated_js_test(
test_code=generated_test_source, function_to_optimize=function_to_optimize, mode=TestingMode.BEHAVIOR
Expand Down
Loading