mirror of
https://github.com/codeflash-ai/codeflash-internal.git
synced 2026-05-04 18:25:18 +00:00
Fix await in non-async describe callbacks
The AI service LLM sometimes generates test code with `await` statements inside describe() callbacks that are not marked as `async`. This causes runtime errors in Vitest/Jest: "Error: `await` is only allowed within async functions" Tree-sitter validation passes because top-level await is valid in ESM, but the describe() callback is not top-level - it's a regular function that needs the `async` keyword. This fix adds a post-processing step that detects describe blocks containing await statements at the describe level (not inside nested test/it functions) and adds the `async` keyword to the callback. Fixes trace ID: caa203f1-58fe-4052-b140-ac0f9c5d6b77 Changes: - Added fix_await_in_describe_callbacks() function to testgen.py - Integrated fix into test generation pipeline after strip_js_extensions() - Added comprehensive unit tests covering single/nested describes, edge cases - All tests pass (10/10 new tests, 23/23 existing tests) Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
Codeflash Bot
parent
0abc6bf1e3
commit
c7f2897a42
2 changed files with 284 additions and 0 deletions
|
|
@ -153,6 +153,125 @@ def strip_js_extensions(source: str) -> str:
|
|||
return _JEST_MOCK_EXTENSION_PATTERN.sub(r"\1\2\4", source)
|
||||
|
||||
|
||||
def fix_await_in_describe_callbacks(source: str) -> str:
|
||||
"""Fix await statements in non-async describe() callbacks.
|
||||
|
||||
The LLM sometimes generates code with `await` statements inside describe() callbacks
|
||||
that are not marked as `async`. This causes runtime errors in Vitest/Jest:
|
||||
"Error: `await` is only allowed within async functions"
|
||||
|
||||
Tree-sitter validation passes because top-level await is valid in ESM, but
|
||||
the describe() callback is not top-level - it's a regular function that needs
|
||||
the `async` keyword.
|
||||
|
||||
This function detects describe blocks containing await statements at the describe
|
||||
level (not inside nested test/it functions) and adds the `async` keyword to
|
||||
the callback. It processes nested describe blocks recursively.
|
||||
|
||||
Args:
|
||||
source: The test source code.
|
||||
|
||||
Returns:
|
||||
Source code with async added to describe callbacks that contain await.
|
||||
|
||||
Example:
|
||||
Input:
|
||||
describe('test', () => {
|
||||
const x = await import('./x.js');
|
||||
test('works', () => { ... });
|
||||
});
|
||||
|
||||
Output:
|
||||
describe('test', async () => {
|
||||
const x = await import('./x.js');
|
||||
test('works', () => { ... });
|
||||
});
|
||||
"""
|
||||
if not source or "describe" not in source or "await" not in source:
|
||||
return source
|
||||
|
||||
# Use a simple regex replacement approach that works line-by-line
|
||||
# Pattern: describe('...', () => { where there's await before next test/it/describe
|
||||
# We'll process multiple times to handle nested describes
|
||||
|
||||
result = source
|
||||
max_iterations = 10 # Prevent infinite loops
|
||||
iteration = 0
|
||||
|
||||
while iteration < max_iterations:
|
||||
iteration += 1
|
||||
changed = False
|
||||
lines = result.split("\n")
|
||||
result_lines = []
|
||||
i = 0
|
||||
|
||||
while i < len(lines):
|
||||
line = lines[i]
|
||||
|
||||
# Check if this line starts a describe block with a non-async arrow function callback
|
||||
# Pattern: describe('...', () => { or describe("...", () => {
|
||||
describe_match = re.match(
|
||||
r'^(\s*describe\s*\(\s*["\'][^"\']*["\']\s*,\s*)(async\s+)?(\(\s*\)\s*=>\s*\{.*)$',
|
||||
line,
|
||||
)
|
||||
|
||||
if describe_match:
|
||||
prefix = describe_match.group(1) # "describe('test', "
|
||||
existing_async = describe_match.group(2) # "async " or None
|
||||
callback = describe_match.group(3) # "() => {"
|
||||
|
||||
# If already async, keep as-is
|
||||
if existing_async:
|
||||
result_lines.append(line)
|
||||
i += 1
|
||||
continue
|
||||
|
||||
# Find the closing brace for this describe block
|
||||
brace_count = line.count("{") - line.count("}")
|
||||
block_lines = [line]
|
||||
i += 1
|
||||
|
||||
# Collect the full describe block
|
||||
while i < len(lines) and brace_count > 0:
|
||||
current_line = lines[i]
|
||||
block_lines.append(current_line)
|
||||
brace_count += current_line.count("{") - current_line.count("}")
|
||||
i += 1
|
||||
|
||||
# Check if there's await at the describe level
|
||||
# (before any nested test/it/describe/beforeEach/afterEach)
|
||||
has_await_at_describe_level = False
|
||||
for block_line in block_lines[1:]: # Skip the describe line itself
|
||||
# If we hit a nested test/it/describe/beforeEach/afterEach, stop checking
|
||||
if re.search(r'\b(test|it|describe|beforeEach|afterEach|beforeAll|afterAll)\s*\(', block_line):
|
||||
break
|
||||
# Check for await statements
|
||||
if re.search(r'\bawait\s+', block_line):
|
||||
has_await_at_describe_level = True
|
||||
break
|
||||
|
||||
# If we found await at describe level, add async to the callback
|
||||
if has_await_at_describe_level:
|
||||
fixed_line = prefix + "async " + callback
|
||||
result_lines.append(fixed_line)
|
||||
result_lines.extend(block_lines[1:])
|
||||
changed = True
|
||||
else:
|
||||
# No await at describe level, keep original
|
||||
result_lines.extend(block_lines)
|
||||
else:
|
||||
result_lines.append(line)
|
||||
i += 1
|
||||
|
||||
result = "\n".join(result_lines)
|
||||
|
||||
# If no changes were made, we're done
|
||||
if not changed:
|
||||
break
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _resolve_import(function_name: str, module_path: str) -> tuple[str, str, str]:
|
||||
"""Determine import style and binding name for a JS/TS function.
|
||||
|
||||
|
|
@ -526,6 +645,9 @@ async def testgen_javascript(
|
|||
# Strip incorrect file extensions from import paths (LLMs sometimes add .js to .ts imports)
|
||||
generated_test_source = strip_js_extensions(generated_test_source)
|
||||
|
||||
# Fix await statements in non-async describe callbacks (LLMs sometimes generate these)
|
||||
generated_test_source = fix_await_in_describe_callbacks(generated_test_source)
|
||||
|
||||
ph(request.user, "aiservice-testgen-tests-generated", properties={"language": language})
|
||||
|
||||
if hasattr(request, "should_log_features") and request.should_log_features:
|
||||
|
|
|
|||
162
django/aiservice/tests/testgen/test_await_in_describe_fix.py
Normal file
162
django/aiservice/tests/testgen/test_await_in_describe_fix.py
Normal file
|
|
@ -0,0 +1,162 @@
|
|||
"""Tests for fixing await statements in non-async describe callbacks.
|
||||
|
||||
The LLM sometimes generates await statements inside describe() callbacks that are not marked as async.
|
||||
This causes runtime errors in Vitest/Jest even though tree-sitter validation passes.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from core.languages.js_ts.testgen import fix_await_in_describe_callbacks
|
||||
|
||||
|
||||
class TestFixAwaitInDescribeCallbacks:
|
||||
"""Tests for fixing await in non-async describe callbacks."""
|
||||
|
||||
def test_fixes_single_await_in_describe(self) -> None:
|
||||
"""Test fixing a single await import in describe callback."""
|
||||
code = """import { describe, test, expect } from 'vitest';
|
||||
|
||||
describe('myFunction', () => {
|
||||
const module = await import('./module.js');
|
||||
|
||||
test('should work', () => {
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});"""
|
||||
|
||||
expected = """import { describe, test, expect } from 'vitest';
|
||||
|
||||
describe('myFunction', async () => {
|
||||
const module = await import('./module.js');
|
||||
|
||||
test('should work', () => {
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});"""
|
||||
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
assert result == expected
|
||||
|
||||
def test_fixes_multiple_awaits_in_describe(self) -> None:
|
||||
"""Test fixing multiple await imports in describe callback."""
|
||||
code = """describe('registerModelsCli', () => {
|
||||
const modelsModule = await import('../commands/models.js');
|
||||
const cliUtilsModule = await import('./cli-utils.js');
|
||||
const runtimeModule = await import('../runtime.js');
|
||||
|
||||
test('should register', () => {
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});"""
|
||||
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
assert "describe('registerModelsCli', async () => {" in result
|
||||
assert "await import('../commands/models.js')" in result
|
||||
|
||||
def test_fixes_nested_describes_with_await(self) -> None:
|
||||
"""Test fixing await in nested describe blocks."""
|
||||
code = """describe('outer', () => {
|
||||
const a = await import('./a.js');
|
||||
|
||||
describe('inner', () => {
|
||||
const b = await import('./b.js');
|
||||
|
||||
test('works', () => {
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});
|
||||
});"""
|
||||
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
# Both describe blocks should be async
|
||||
assert result.count("describe('outer', async () => {") == 1
|
||||
assert result.count("describe('inner', async () => {") == 1
|
||||
|
||||
def test_does_not_modify_already_async_describe(self) -> None:
|
||||
"""Test that already async describe blocks are not modified."""
|
||||
code = """describe('myFunction', async () => {
|
||||
const module = await import('./module.js');
|
||||
|
||||
test('should work', () => {
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});"""
|
||||
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
# Should remain unchanged
|
||||
assert result == code
|
||||
|
||||
def test_does_not_modify_describe_without_await(self) -> None:
|
||||
"""Test that describe blocks without await are not modified."""
|
||||
code = """describe('myFunction', () => {
|
||||
const module = require('./module.js');
|
||||
|
||||
test('should work', () => {
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});"""
|
||||
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
# Should remain unchanged
|
||||
assert result == code
|
||||
|
||||
def test_handles_await_in_test_function(self) -> None:
|
||||
"""Test that await inside test functions is not affected."""
|
||||
code = """describe('myFunction', () => {
|
||||
test('should work', async () => {
|
||||
const result = await myAsyncFunction();
|
||||
expect(result).toBe(true);
|
||||
});
|
||||
});"""
|
||||
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
# describe should remain non-async since there's no await at describe level
|
||||
assert "describe('myFunction', () => {" in result
|
||||
|
||||
def test_fixes_describe_with_await_at_top_level(self) -> None:
|
||||
"""Test fixing describe with await before any test definitions."""
|
||||
code = """describe('myFunction', () => {
|
||||
// Import modules at top
|
||||
const modelsModule = await import('../commands/models.js');
|
||||
const cliUtilsModule = await import('./cli-utils.js');
|
||||
const runtimeModule = await import('../runtime.js');
|
||||
|
||||
let program;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
program = new CommandStub('root');
|
||||
});
|
||||
|
||||
test('basic test', () => {
|
||||
expect(true).toBe(true);
|
||||
});
|
||||
});"""
|
||||
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
assert "describe('myFunction', async () => {" in result
|
||||
|
||||
def test_handles_various_callback_styles(self) -> None:
|
||||
"""Test handling different arrow function styles."""
|
||||
# Arrow function with space
|
||||
code1 = "describe('test', () => {\n const x = await import('./x.js');\n});"
|
||||
result1 = fix_await_in_describe_callbacks(code1)
|
||||
assert "describe('test', async () => {" in result1
|
||||
|
||||
# Arrow function without space (edge case)
|
||||
code2 = "describe('test', ()=>{\n const x = await import('./x.js');\n});"
|
||||
result2 = fix_await_in_describe_callbacks(code2)
|
||||
assert "async" in result2
|
||||
|
||||
def test_empty_code_returns_unchanged(self) -> None:
|
||||
"""Test that empty code is returned unchanged."""
|
||||
assert fix_await_in_describe_callbacks("") == ""
|
||||
|
||||
def test_code_without_describe_returns_unchanged(self) -> None:
|
||||
"""Test that code without describe blocks is returned unchanged."""
|
||||
code = """const x = 5;
|
||||
test('works', () => {
|
||||
expect(x).toBe(5);
|
||||
});"""
|
||||
result = fix_await_in_describe_callbacks(code)
|
||||
assert result == code
|
||||
Loading…
Reference in a new issue