mirror of
https://github.com/codeflash-ai/codeflash-agent.git
synced 2026-05-04 18:25:19 +00:00
2c9f2ad8de
Restore the old InjectPerfOnly behavior where call-site identifiers are the source line number of the instrumented statement. Also fix the sync integration test to properly apply the decorator and write the helper file, and remove dead imports from test_instrumentation.
771 lines
23 KiB
Python
771 lines
23 KiB
Python
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from codeflash_python._model import FunctionParent, TestingMode
|
|
from codeflash_python.analysis._discovery import FunctionToOptimize
|
|
from codeflash_python.test_discovery.models import CodePosition
|
|
from codeflash_python.testing._instrument_async import (
|
|
ASYNC_HELPER_FILENAME,
|
|
add_async_decorator_to_function,
|
|
get_decorator_name_for_mode,
|
|
)
|
|
from codeflash_python.testing._instrumentation import (
|
|
inject_profiling_into_existing_test,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_dir():
|
|
"""Create a temporary directory for test files."""
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
yield Path(temp)
|
|
|
|
|
|
# @pytest.fixture
|
|
# def unique_test_iteration():
|
|
# """Provide a unique test iteration ID and clean up database after test."""
|
|
# # Generate unique iteration ID
|
|
# iteration_id = str(uuid.uuid4())[:8]
|
|
|
|
# # Store original environment variable
|
|
# original_iteration = os.environ.get("CODEFLASH_TEST_ITERATION")
|
|
|
|
# # Set unique iteration for this test
|
|
# os.environ["CODEFLASH_TEST_ITERATION"] = iteration_id
|
|
|
|
# try:
|
|
# yield iteration_id
|
|
# finally:
|
|
# # Cleanup: restore original environment and delete database file
|
|
# if original_iteration is not None:
|
|
# os.environ["CODEFLASH_TEST_ITERATION"] = original_iteration
|
|
# elif "CODEFLASH_TEST_ITERATION" in os.environ:
|
|
# del os.environ["CODEFLASH_TEST_ITERATION"]
|
|
|
|
# # Clean up database file
|
|
# try:
|
|
# from codeflash.code_utils.codeflash_wrap_decorator import get_run_tmp_file
|
|
|
|
# db_path = get_run_tmp_file(Path(f"test_return_values_{iteration_id}.sqlite"))
|
|
# if db_path.exists():
|
|
# db_path.unlink()
|
|
# except Exception:
|
|
# pass # Ignore cleanup errors
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_async_decorator_application_behavior_mode(temp_dir):
|
|
async_function_code = '''
|
|
import asyncio
|
|
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Simple async function for testing."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(async_function_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
file_path=test_file,
|
|
parents=[],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert decorator_added
|
|
modified_code = test_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR)
|
|
code_with_decorator = async_function_code.replace(
|
|
"async def async_function",
|
|
f"@{decorator_name}\nasync def async_function",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert modified_code.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_async_decorator_application_performance_mode(temp_dir):
|
|
async_function_code = '''
|
|
import asyncio
|
|
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Simple async function for testing."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(async_function_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
file_path=test_file,
|
|
parents=[],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, func, TestingMode.PERFORMANCE
|
|
)
|
|
|
|
assert decorator_added
|
|
modified_code = test_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.PERFORMANCE)
|
|
code_with_decorator = async_function_code.replace(
|
|
"async def async_function",
|
|
f"@{decorator_name}\nasync def async_function",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert modified_code.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_async_decorator_application_concurrency_mode(temp_dir):
|
|
"""Test that CONCURRENCY mode applies the codeflash_concurrency_async decorator."""
|
|
async_function_code = '''
|
|
import asyncio
|
|
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Simple async function for testing."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(async_function_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
file_path=test_file,
|
|
parents=[],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, func, TestingMode.CONCURRENCY
|
|
)
|
|
|
|
assert decorator_added
|
|
modified_code = test_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.CONCURRENCY)
|
|
code_with_decorator = async_function_code.replace(
|
|
"async def async_function",
|
|
f"@{decorator_name}\nasync def async_function",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert modified_code.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_async_class_method_decorator_application(temp_dir):
|
|
async_class_code = '''
|
|
import asyncio
|
|
|
|
class Calculator:
|
|
"""Test class with async methods."""
|
|
|
|
async def async_method(self, a: int, b: int) -> int:
|
|
"""Async method in class."""
|
|
await asyncio.sleep(0.005)
|
|
return a ** b
|
|
|
|
def sync_method(self, a: int, b: int) -> int:
|
|
"""Sync method in class."""
|
|
return a - b
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(async_class_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_method",
|
|
file_path=test_file,
|
|
parents=[FunctionParent(name="Calculator", type="ClassDef")],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert decorator_added
|
|
modified_code = test_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR)
|
|
code_with_decorator = async_class_code.replace(
|
|
" async def async_method",
|
|
f" @{decorator_name}\n async def async_method",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert modified_code.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_async_decorator_no_duplicate_application(temp_dir):
|
|
# Case 1: Old-style import already present — injector should detect and skip
|
|
already_decorated_code = '''
|
|
from codeflash_python.runtime._codeflash_wrap_decorator import codeflash_behavior_async
|
|
import asyncio
|
|
|
|
@codeflash_behavior_async
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Already decorated async function."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(already_decorated_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
file_path=test_file,
|
|
parents=[],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
# Should not add duplicate decorator
|
|
assert not decorator_added
|
|
|
|
# Case 2: Inline definition already present — injector should detect and skip
|
|
already_inline_code = '''
|
|
import asyncio
|
|
|
|
def codeflash_behavior_async(func):
|
|
return func
|
|
|
|
@codeflash_behavior_async
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Already decorated async function."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
test_file2 = temp_dir / "test_async2.py"
|
|
test_file2.write_text(already_inline_code)
|
|
|
|
func2 = FunctionToOptimize(
|
|
function_name="async_function",
|
|
file_path=test_file2,
|
|
parents=[],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added2, _ = add_async_decorator_to_function(
|
|
test_file2, func2, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
# Should not add duplicate decorator
|
|
assert not decorator_added2
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_inject_profiling_async_function_behavior_mode(temp_dir):
|
|
source_module_code = '''
|
|
import asyncio
|
|
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Simple async function for testing."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
source_file = temp_dir / "my_module.py"
|
|
source_file.write_text(source_module_code)
|
|
|
|
async_test_code = '''
|
|
import asyncio
|
|
import pytest
|
|
from my_module import async_function
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_function():
|
|
"""Test async function behavior."""
|
|
result = await async_function(5, 3)
|
|
assert result == 15
|
|
|
|
result2 = await async_function(2, 4)
|
|
assert result2 == 8
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(async_test_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
parents=[],
|
|
file_path=Path("my_module.py"),
|
|
is_async=True,
|
|
)
|
|
|
|
# First instrument the source module
|
|
from codeflash_python.testing._instrument_async import (
|
|
add_async_decorator_to_function,
|
|
)
|
|
|
|
source_success, _ = add_async_decorator_to_function(
|
|
source_file, func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert source_success is True
|
|
|
|
# Verify the file was modified with exact expected output
|
|
instrumented_source = source_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR)
|
|
code_with_decorator = source_module_code.replace(
|
|
"async def async_function",
|
|
f"@{decorator_name}\nasync def async_function",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert instrumented_source.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
success, instrumented_test_code = inject_profiling_into_existing_test(
|
|
test_file,
|
|
[CodePosition(8, 18), CodePosition(11, 19)],
|
|
func,
|
|
temp_dir,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
|
|
# For async functions, once source is decorated, test injection should fail
|
|
# This is expected behavior - async instrumentation happens at the decorator level
|
|
assert success is False
|
|
assert instrumented_test_code is None
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_inject_profiling_async_function_performance_mode(temp_dir):
|
|
source_module_code = '''
|
|
import asyncio
|
|
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Simple async function for testing."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
source_file = temp_dir / "my_module.py"
|
|
source_file.write_text(source_module_code)
|
|
|
|
# Create the test file
|
|
async_test_code = '''
|
|
import asyncio
|
|
import pytest
|
|
from my_module import async_function
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_function():
|
|
"""Test async function performance."""
|
|
result = await async_function(5, 3)
|
|
assert result == 15
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(async_test_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
parents=[],
|
|
file_path=Path("my_module.py"),
|
|
is_async=True,
|
|
)
|
|
|
|
# First instrument the source module
|
|
from codeflash_python.testing._instrument_async import (
|
|
add_async_decorator_to_function,
|
|
)
|
|
|
|
source_success, _ = add_async_decorator_to_function(
|
|
source_file, func, TestingMode.PERFORMANCE
|
|
)
|
|
|
|
assert source_success is True
|
|
|
|
# Verify the file was modified with exact expected output
|
|
instrumented_source = source_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.PERFORMANCE)
|
|
code_with_decorator = source_module_code.replace(
|
|
"async def async_function",
|
|
f"@{decorator_name}\nasync def async_function",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert instrumented_source.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
# Now test the full pipeline with source module path
|
|
success, instrumented_test_code = inject_profiling_into_existing_test(
|
|
test_file,
|
|
[CodePosition(8, 18)],
|
|
func,
|
|
temp_dir,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
|
|
# For async functions, once source is decorated, test injection should fail
|
|
# This is expected behavior - async instrumentation happens at the decorator level
|
|
assert success is False
|
|
assert instrumented_test_code is None
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_mixed_sync_async_instrumentation(temp_dir):
|
|
source_module_code = '''
|
|
import asyncio
|
|
|
|
def sync_function(x: int, y: int) -> int:
|
|
"""Regular sync function."""
|
|
return x * y
|
|
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Simple async function."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
source_file = temp_dir / "my_module.py"
|
|
source_file.write_text(source_module_code)
|
|
|
|
mixed_test_code = '''
|
|
import asyncio
|
|
import pytest
|
|
from my_module import sync_function, async_function
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mixed_functions():
|
|
"""Test both sync and async functions."""
|
|
sync_result = sync_function(10, 5)
|
|
assert sync_result == 50
|
|
|
|
async_result = await async_function(3, 4)
|
|
assert async_result == 12
|
|
'''
|
|
|
|
test_file = temp_dir / "test_mixed.py"
|
|
test_file.write_text(mixed_test_code)
|
|
|
|
async_func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
parents=[],
|
|
file_path=Path("my_module.py"),
|
|
is_async=True,
|
|
)
|
|
|
|
from codeflash_python.testing._instrument_async import (
|
|
add_async_decorator_to_function,
|
|
)
|
|
|
|
source_success, _ = add_async_decorator_to_function(
|
|
source_file, async_func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert source_success
|
|
|
|
# Verify the file was modified
|
|
instrumented_source = source_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR)
|
|
code_with_decorator = source_module_code.replace(
|
|
"async def async_function",
|
|
f"@{decorator_name}\nasync def async_function",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert instrumented_source.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
success, instrumented_test_code = inject_profiling_into_existing_test(
|
|
test_file,
|
|
[CodePosition(8, 18), CodePosition(11, 19)],
|
|
async_func,
|
|
temp_dir,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
|
|
# Async functions should not be instrumented at the test level
|
|
assert not success
|
|
assert instrumented_test_code is None
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_async_function_qualified_name_handling(temp_dir):
|
|
nested_async_code = '''
|
|
import asyncio
|
|
|
|
class OuterClass:
|
|
class InnerClass:
|
|
async def nested_async_method(self, x: int) -> int:
|
|
"""Nested async method."""
|
|
await asyncio.sleep(0.001)
|
|
return x * 2
|
|
'''
|
|
|
|
test_file = temp_dir / "test_nested.py"
|
|
test_file.write_text(nested_async_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="nested_async_method",
|
|
file_path=test_file,
|
|
parents=[
|
|
FunctionParent(name="OuterClass", type="ClassDef"),
|
|
FunctionParent(name="InnerClass", type="ClassDef"),
|
|
],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert decorator_added
|
|
modified_code = test_file.read_text()
|
|
from codeflash_python.analysis._formatter import sort_imports
|
|
|
|
decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR)
|
|
code_with_decorator = nested_async_code.replace(
|
|
" async def nested_async_method",
|
|
f" @{decorator_name}\n async def nested_async_method",
|
|
)
|
|
code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}"
|
|
expected = sort_imports(code=code_with_import, float_to_top=True)
|
|
assert modified_code.strip() == expected.strip()
|
|
assert (temp_dir / ASYNC_HELPER_FILENAME).exists()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_async_decorator_with_existing_decorators(temp_dir):
|
|
"""Test async decorator application when function already has other decorators."""
|
|
decorated_async_code = '''
|
|
import asyncio
|
|
from functools import wraps
|
|
|
|
def my_decorator(func):
|
|
@wraps(func)
|
|
async def wrapper(*args, **kwargs):
|
|
return await func(*args, **kwargs)
|
|
return wrapper
|
|
|
|
@my_decorator
|
|
async def async_function(x: int, y: int) -> int:
|
|
"""Async function with existing decorator."""
|
|
await asyncio.sleep(0.01)
|
|
return x * y
|
|
'''
|
|
|
|
test_file = temp_dir / "test_async.py"
|
|
test_file.write_text(decorated_async_code)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_function",
|
|
file_path=test_file,
|
|
parents=[],
|
|
is_async=True,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert decorator_added
|
|
modified_code = test_file.read_text()
|
|
# Should add codeflash decorator above existing decorators
|
|
assert "@codeflash_behavior_async" in modified_code
|
|
assert "@my_decorator" in modified_code
|
|
# Codeflash decorator should come first
|
|
codeflash_pos = modified_code.find("@codeflash_behavior_async")
|
|
my_decorator_pos = modified_code.find("@my_decorator")
|
|
assert codeflash_pos < my_decorator_pos
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_sync_function_not_affected_by_async_logic(temp_dir):
|
|
sync_function_code = '''
|
|
def sync_function(x: int, y: int) -> int:
|
|
"""Regular sync function."""
|
|
return x + y
|
|
'''
|
|
|
|
test_file = temp_dir / "test_sync.py"
|
|
test_file.write_text(sync_function_code)
|
|
|
|
sync_func = FunctionToOptimize(
|
|
function_name="sync_function",
|
|
file_path=test_file,
|
|
parents=[],
|
|
is_async=False,
|
|
)
|
|
|
|
decorator_added, _ = add_async_decorator_to_function(
|
|
test_file, sync_func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert not decorator_added
|
|
# File should not be modified for sync functions
|
|
modified_code = test_file.read_text()
|
|
assert modified_code == sync_function_code
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32", reason="pending support for asyncio on windows"
|
|
)
|
|
def test_inject_profiling_async_multiple_calls_same_test(temp_dir):
|
|
"""Test that multiple async function calls within the same test function get correctly numbered 0, 1, 2, etc."""
|
|
source_module_code = '''
|
|
import asyncio
|
|
|
|
async def async_sorter(items):
|
|
"""Simple async sorter for testing."""
|
|
await asyncio.sleep(0.001)
|
|
return sorted(items)
|
|
'''
|
|
|
|
source_file = temp_dir / "async_sorter.py"
|
|
source_file.write_text(source_module_code)
|
|
|
|
test_code_multiple_calls = """
|
|
import asyncio
|
|
import pytest
|
|
from async_sorter import async_sorter
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_call():
|
|
result = await async_sorter([42])
|
|
assert result == [42]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_calls():
|
|
result1 = await async_sorter([3, 1, 2])
|
|
result2 = await async_sorter([5, 4])
|
|
result3 = await async_sorter([9, 8, 7, 6])
|
|
assert result1 == [1, 2, 3]
|
|
assert result2 == [4, 5]
|
|
assert result3 == [6, 7, 8, 9]
|
|
"""
|
|
|
|
test_file = temp_dir / "test_async_sorter.py"
|
|
test_file.write_text(test_code_multiple_calls)
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="async_sorter",
|
|
parents=[],
|
|
file_path=Path("async_sorter.py"),
|
|
is_async=True,
|
|
)
|
|
|
|
# First instrument the source module with async decorators
|
|
from codeflash_python.testing._instrument_async import (
|
|
add_async_decorator_to_function,
|
|
)
|
|
|
|
source_success, _ = add_async_decorator_to_function(
|
|
source_file, func, TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert source_success
|
|
|
|
# Verify the file was modified
|
|
instrumented_source = source_file.read_text()
|
|
assert "@codeflash_behavior_async" in instrumented_source
|
|
|
|
import ast
|
|
|
|
tree = ast.parse(test_code_multiple_calls)
|
|
call_positions = []
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Await) and isinstance(node.value, ast.Call):
|
|
if (
|
|
hasattr(node.value.func, "id")
|
|
and node.value.func.id == "async_sorter"
|
|
) or (
|
|
hasattr(node.value.func, "attr")
|
|
and node.value.func.attr == "async_sorter"
|
|
):
|
|
call_positions.append(
|
|
CodePosition(node.lineno, node.col_offset)
|
|
)
|
|
|
|
assert len(call_positions) == 4
|
|
|
|
success, instrumented_test_code = inject_profiling_into_existing_test(
|
|
test_file, call_positions, func, temp_dir, mode=TestingMode.BEHAVIOR
|
|
)
|
|
|
|
assert success
|
|
assert instrumented_test_code is not None
|
|
|
|
assert "_codeflash_call_site.set('8')" in instrumented_test_code
|
|
assert "_codeflash_call_site.set('13')" in instrumented_test_code
|
|
assert "_codeflash_call_site.set('14')" in instrumented_test_code
|
|
assert "_codeflash_call_site.set('15')" in instrumented_test_code
|
|
|
|
assert 1 == instrumented_test_code.count(
|
|
"_codeflash_call_site.set('8')"
|
|
)
|
|
assert 1 == instrumented_test_code.count(
|
|
"_codeflash_call_site.set('13')"
|
|
)
|
|
assert 1 == instrumented_test_code.count(
|
|
"_codeflash_call_site.set('14')"
|
|
)
|
|
assert 1 == instrumented_test_code.count(
|
|
"_codeflash_call_site.set('15')"
|
|
)
|
|
|
|
|