test: sync dual-changed test files from main with omni-java fixes

Updates inject_profiling_into_existing_test calls to include test_string
parameter. Takes main's test refactoring for multi-file code replacement
and codeflash capture.
This commit is contained in:
Kevin Turcios 2026-03-02 15:30:16 -05:00
parent 19bd6e4bad
commit bd3ec8f09d
4 changed files with 116 additions and 197 deletions

View file

@ -141,24 +141,22 @@ def run_codeflash_command(
def build_command(
cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path, benchmarks_root: pathlib.Path | None = None
) -> list[str]:
repo_root = pathlib.Path(__file__).resolve().parent.parent.parent
python_path = os.path.relpath(repo_root / "codeflash" / "main.py", cwd)
python_path = "../../../codeflash/main.py" if "code_directories" in str(cwd) else "../codeflash/main.py"
base_command = ["uv", "run", "--no-project", python_path, "--file", config.file_path, "--no-pr"]
if config.function_name:
base_command.extend(["--function", config.function_name])
# Check if codeflash config exists (pyproject.toml or codeflash.toml) - if so, don't override it
has_codeflash_config = (cwd / "codeflash.toml").exists()
if not has_codeflash_config:
# Check if pyproject.toml exists with codeflash config - if so, don't override it
pyproject_path = cwd / "pyproject.toml"
has_codeflash_config = False
if pyproject_path.exists():
with contextlib.suppress(Exception), open(pyproject_path, "rb") as f:
pyproject_data = tomllib.load(f)
has_codeflash_config = "tool" in pyproject_data and "codeflash" in pyproject_data["tool"]
# Only pass --tests-root and --module-root if they're not configured
# Only pass --tests-root and --module-root if they're not configured in pyproject.toml
if not has_codeflash_config:
base_command.extend(["--tests-root", str(test_root), "--module-root", str(cwd)])
@ -222,7 +220,7 @@ def validate_output(stdout: str, return_code: int, expected_improvement_pct: int
return False
if config.expected_unit_test_files is not None:
# Match the per-function test discovery message from function_optimizer.py
# Match the per-function discovery message from function_optimizer.py
# Format: "Discovered X existing unit test files, Y replay test files, and Z concolic..."
unit_test_files_match = re.search(r"Discovered (\d+) existing unit test files?", stdout)
if not unit_test_files_match:

View file

@ -7,8 +7,8 @@ from pathlib import Path
from codeflash.code_utils.code_utils import get_run_tmp_file
from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer
from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode, TestType, VerificationType
from codeflash.optimization.function_optimizer import FunctionOptimizer
from codeflash.verification.equivalence import compare_test_results
from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
from codeflash.verification.test_runner import execute_test_subprocess
@ -459,7 +459,7 @@ class MyClass:
file_path=sample_code_path,
parents=[FunctionParent(name="MyClass", type="ClassDef")],
)
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer.test_files = TestFiles(
test_files=[
TestFile(
@ -475,8 +475,8 @@ class MyClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert len(test_results) == 3
@ -508,8 +508,8 @@ class MyClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
match, _ = compare_test_results(test_results, test_results2)
@ -582,7 +582,7 @@ class MyClass(ParentClass):
file_path=sample_code_path,
parents=[FunctionParent(name="MyClass", type="ClassDef")],
)
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer.test_files = TestFiles(
test_files=[
TestFile(
@ -598,8 +598,8 @@ class MyClass(ParentClass):
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert len(test_results) == 3
@ -632,8 +632,8 @@ class MyClass(ParentClass):
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
@ -709,7 +709,7 @@ class MyClass:
file_path=sample_code_path,
parents=[FunctionParent(name="MyClass", type="ClassDef")],
)
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer.test_files = TestFiles(
test_files=[
TestFile(
@ -725,8 +725,8 @@ class MyClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
@ -761,8 +761,8 @@ class MyClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
@ -872,7 +872,7 @@ class AnotherHelperClass:
file_path=fto_file_path,
parents=[FunctionParent(name="MyClass", type="ClassDef")],
)
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer.test_files = TestFiles(
test_files=[
TestFile(
@ -889,8 +889,8 @@ class AnotherHelperClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
@ -910,8 +910,8 @@ class AnotherHelperClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
@ -1021,7 +1021,7 @@ class AnotherHelperClass:
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer.test_files = TestFiles(
test_files=[
TestFile(
@ -1049,13 +1049,13 @@ class AnotherHelperClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
# Remove instrumentation
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
assert len(test_results.test_results) == 4
assert test_results[0].id.test_function_name == "test_helper_classes"
@ -1101,12 +1101,12 @@ class MyClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
# Remove instrumentation
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
# Now, this fto_code mutates the instance so it should fail
mutated_fto_code = """
@ -1140,12 +1140,12 @@ class MyClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
# Remove instrumentation
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
match, _ = compare_test_results(test_results, mutated_test_results)
assert not match
@ -1179,12 +1179,12 @@ class MyClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
# Remove instrumentation
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
match, _ = compare_test_results(test_results, no_helper1_test_results)
assert match
@ -1446,7 +1446,7 @@ def calculate_portfolio_metrics(
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer.test_files = TestFiles(
test_files=[
TestFile(
@ -1471,13 +1471,13 @@ def calculate_portfolio_metrics(
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
# Remove instrumentation
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
# Now, let's say we optimize the code and make changes.
new_fto_code = """import math
@ -1538,12 +1538,12 @@ def calculate_portfolio_metrics(
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
# Remove instrumentation
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
matched, diffs = compare_test_results(test_results, modified_test_results)
assert not matched
@ -1601,12 +1601,12 @@ def calculate_portfolio_metrics(
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
# Remove instrumentation
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
matched, diffs = compare_test_results(test_results, modified_test_results_2)
# now the test should match and no diffs should be found
assert len(diffs) == 0
@ -1671,7 +1671,7 @@ class SlotsClass:
file_path=sample_code_path,
parents=[FunctionParent(name="SlotsClass", type="ClassDef")],
)
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
func_optimizer.test_files = TestFiles(
test_files=[
TestFile(
@ -1687,8 +1687,8 @@ class SlotsClass:
test_env=test_env,
test_files=func_optimizer.test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)

View file

@ -15,8 +15,9 @@ from codeflash.code_utils.instrument_existing_tests import (
FunctionImportedAsVisitor,
inject_profiling_into_existing_test,
)
from codeflash.languages.python.static_analysis.line_profile_utils import add_decorator_imports
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer
from codeflash.languages.python.static_analysis.line_profile_utils import add_decorator_imports
from codeflash.models.models import (
CodeOptimizationContext,
CodePosition,
@ -27,7 +28,6 @@ from codeflash.models.models import (
TestsInFile,
TestType,
)
from codeflash.optimization.function_optimizer import FunctionOptimizer
from codeflash.verification.verification_utils import TestConfig
codeflash_wrap_string = """def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, codeflash_cur, codeflash_con, *args, **kwargs):
@ -434,7 +434,7 @@ def test_sort():
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_env = os.environ.copy()
test_env["CODEFLASH_TEST_ITERATION"] = "0"
test_env["CODEFLASH_LOOP_INDEX"] = "1"
@ -454,8 +454,8 @@ def test_sort():
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -489,8 +489,8 @@ def test_sort():
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results_perf[0].id.function_getting_tested == "sorter"
@ -541,8 +541,8 @@ result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
line_profiler_output_file=line_profiler_output_file,
)
@ -695,14 +695,14 @@ def test_sort_parametrized(input, expected_output):
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
testing_type=TestingMode.BEHAVIOR,
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -755,8 +755,8 @@ result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results_perf[0].id.function_getting_tested == "sorter"
@ -812,8 +812,8 @@ result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
line_profiler_output_file=line_profiler_output_file,
)
@ -984,14 +984,14 @@ def test_sort_parametrized_loop(input, expected_output):
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
testing_type=TestingMode.BEHAVIOR,
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -1081,8 +1081,8 @@ result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 2
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
@ -1171,8 +1171,8 @@ result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 2
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
line_profiler_output_file=line_profiler_output_file,
)
@ -1341,14 +1341,14 @@ def test_sort():
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
testing_type=TestingMode.BEHAVIOR,
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -1389,8 +1389,8 @@ def test_sort():
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -1453,8 +1453,8 @@ result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 2
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
line_profiler_output_file=line_profiler_output_file,
)
@ -1723,14 +1723,14 @@ class TestPigLatin(unittest.TestCase):
test_framework="unittest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
testing_type=TestingMode.BEHAVIOR,
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -1779,8 +1779,8 @@ result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -1973,14 +1973,14 @@ import unittest
test_framework="unittest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
testing_type=TestingMode.BEHAVIOR,
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -2034,8 +2034,8 @@ result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 2
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -2229,14 +2229,14 @@ import unittest
test_framework="unittest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
test_env=test_env,
testing_type=TestingMode.BEHAVIOR,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -2290,8 +2290,8 @@ result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 2
testing_type=TestingMode.PERFORMANCE,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -2481,14 +2481,14 @@ import unittest
test_framework="unittest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=f, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=f, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
testing_type=TestingMode.BEHAVIOR,
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -2574,8 +2574,8 @@ result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)
assert test_results[0].id.function_getting_tested == "sorter"
@ -3144,7 +3144,7 @@ def test_sleepfunc_sequence_short(n, expected_total_sleep_time):
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_files = TestFiles(
test_files=[
TestFile(
@ -3160,8 +3160,8 @@ def test_sleepfunc_sequence_short(n, expected_total_sleep_time):
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=2,
max_outer_loops=2,
pytest_min_loops=2,
pytest_max_loops=2,
testing_time=0.1,
)
@ -3279,14 +3279,14 @@ import unittest
test_framework="unittest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
test_results, coverage_data = func_optimizer.run_and_parse_tests(
testing_type=TestingMode.PERFORMANCE,
test_env=test_env,
test_files=test_files,
optimization_iteration=0,
min_outer_loops=1,
max_outer_loops=1,
pytest_min_loops=1,
pytest_max_loops=1,
testing_time=0.1,
)

View file

@ -1,8 +1,8 @@
from pathlib import Path
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown, FunctionParent
from codeflash.optimization.function_optimizer import FunctionOptimizer
from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer
from codeflash.models.models import CodeOptimizationContext, CodeStringsMarkdown
from codeflash.verification.verification_utils import TestConfig
@ -106,7 +106,7 @@ def _get_string_usage(text: str) -> Usage:
test_framework="pytest",
pytest_cmd="pytest",
)
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
code_context: CodeOptimizationContext = func_optimizer.get_code_optimization_context().unwrap()
original_helper_code: dict[Path, str] = {}
@ -165,82 +165,3 @@ def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int:
assert new_code.rstrip() == original_main.rstrip() # No Change
assert new_helper_code.rstrip() == expected_helper.rstrip()
def test_optimized_code_for_different_file_not_applied_to_current_file() -> None:
"""Test that optimized code for one file is not incorrectly applied to a different file.
This reproduces the bug from PR #1309 where optimized code for `formatter.py`
was incorrectly applied to `support.py`, causing `normalize_java_code` to be
duplicated. The bug was in `get_optimized_code_for_module` which had a fallback
that applied a single code block to ANY file being processed.
The scenario:
1. `support.py` imports `normalize_java_code` from `formatter.py`
2. AI returns optimized code with a single code block for `formatter.py`
3. BUG: When processing `support.py`, the fallback applies `formatter.py`'s code
4. EXPECTED: No code should be applied to `support.py` since the paths don't match
"""
from codeflash.languages.python.static_analysis.code_extractor import find_preexisting_objects
from codeflash.languages.python.static_analysis.code_replacer import replace_function_definitions_in_module
from codeflash.models.models import CodeStringsMarkdown
root_dir = Path(__file__).parent.parent.resolve()
# Create support.py - the file that imports the helper
support_file = (root_dir / "code_to_optimize/temp_pr1309_support.py").resolve()
original_support = '''from temp_pr1309_formatter import normalize_java_code
class JavaSupport:
"""Support class for Java operations."""
def normalize_code(self, source: str) -> str:
"""Normalize code for deduplication."""
return normalize_java_code(source)
'''
support_file.write_text(original_support, encoding="utf-8")
# AI returns optimized code for formatter.py ONLY (with explicit path)
# This simulates what happens when the AI optimizes the helper function
optimized_markdown = '''```python:code_to_optimize/temp_pr1309_formatter.py
def normalize_java_code(source: str) -> str:
"""Optimized version with fast-path."""
if not source:
return ""
return "\\n".join(line.strip() for line in source.splitlines() if line.strip())
```
'''
preexisting_objects = find_preexisting_objects(original_support)
# Process support.py with the optimized code that's meant for formatter.py
replace_function_definitions_in_module(
function_names=["JavaSupport.normalize_code"],
optimized_code=CodeStringsMarkdown.parse_markdown_code(optimized_markdown),
module_abspath=support_file,
preexisting_objects=preexisting_objects,
project_root_path=root_dir,
)
new_support_code = support_file.read_text(encoding="utf-8")
# Cleanup
support_file.unlink(missing_ok=True)
# CRITICAL: support.py should NOT have normalize_java_code defined!
# The optimized code was for formatter.py, not support.py.
def_count = new_support_code.count("def normalize_java_code")
assert def_count == 0, (
f"Bug: normalize_java_code was incorrectly added to support.py!\n"
f"Found {def_count} definition(s) when there should be 0.\n"
f"The optimized code was for formatter.py, not support.py.\n"
f"Resulting code:\n{new_support_code}"
)
# The file should remain unchanged since no code matched its path
assert new_support_code.strip() == original_support.strip(), (
f"support.py was modified when it shouldn't have been.\n"
f"Original:\n{original_support}\n"
f"New:\n{new_support_code}"
)