Skip to content

Commit a55841b

Browse files
committed
fix: use PythonFunctionOptimizer in tests that depend on Python-specific hooks
1 parent 96c1459 commit a55841b

2 files changed

Lines changed: 26 additions & 26 deletions

File tree

tests/test_codeflash_capture.py

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
99
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
1010
from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode, TestType, VerificationType
11-
from codeflash.optimization.function_optimizer import FunctionOptimizer
11+
from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer
1212
from codeflash.verification.equivalence import compare_test_results
1313
from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
1414
from codeflash.verification.test_runner import execute_test_subprocess
@@ -459,7 +459,7 @@ def __init__(self, x=2):
459459
file_path=sample_code_path,
460460
parents=[FunctionParent(name="MyClass", type="ClassDef")],
461461
)
462-
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
462+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
463463
func_optimizer.test_files = TestFiles(
464464
test_files=[
465465
TestFile(
@@ -582,7 +582,7 @@ def __init__(self, *args, **kwargs):
582582
file_path=sample_code_path,
583583
parents=[FunctionParent(name="MyClass", type="ClassDef")],
584584
)
585-
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
585+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
586586
func_optimizer.test_files = TestFiles(
587587
test_files=[
588588
TestFile(
@@ -709,7 +709,7 @@ def __init__(self, x=2):
709709
file_path=sample_code_path,
710710
parents=[FunctionParent(name="MyClass", type="ClassDef")],
711711
)
712-
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
712+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
713713
func_optimizer.test_files = TestFiles(
714714
test_files=[
715715
TestFile(
@@ -872,7 +872,7 @@ def another_helper(self):
872872
file_path=fto_file_path,
873873
parents=[FunctionParent(name="MyClass", type="ClassDef")],
874874
)
875-
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
875+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
876876
func_optimizer.test_files = TestFiles(
877877
test_files=[
878878
TestFile(
@@ -1021,7 +1021,7 @@ def another_helper(self):
10211021
test_framework="pytest",
10221022
pytest_cmd="pytest",
10231023
)
1024-
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
1024+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
10251025
func_optimizer.test_files = TestFiles(
10261026
test_files=[
10271027
TestFile(
@@ -1055,7 +1055,7 @@ def another_helper(self):
10551055
)
10561056

10571057
# Remove instrumentation
1058-
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
1058+
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
10591059

10601060
assert len(test_results.test_results) == 4
10611061
assert test_results[0].id.test_function_name == "test_helper_classes"
@@ -1106,7 +1106,7 @@ def target_function(self):
11061106
testing_time=0.1,
11071107
)
11081108
# Remove instrumentation
1109-
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
1109+
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
11101110

11111111
# Now, this fto_code mutates the instance so it should fail
11121112
mutated_fto_code = """
@@ -1145,7 +1145,7 @@ def target_function(self):
11451145
testing_time=0.1,
11461146
)
11471147
# Remove instrumentation
1148-
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
1148+
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
11491149
match, _ = compare_test_results(test_results, mutated_test_results)
11501150
assert not match
11511151

@@ -1184,7 +1184,7 @@ def target_function(self):
11841184
testing_time=0.1,
11851185
)
11861186
# Remove instrumentation
1187-
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
1187+
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
11881188
match, _ = compare_test_results(test_results, no_helper1_test_results)
11891189
assert match
11901190

@@ -1446,7 +1446,7 @@ def calculate_portfolio_metrics(
14461446
test_framework="pytest",
14471447
pytest_cmd="pytest",
14481448
)
1449-
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
1449+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
14501450
func_optimizer.test_files = TestFiles(
14511451
test_files=[
14521452
TestFile(
@@ -1477,7 +1477,7 @@ def calculate_portfolio_metrics(
14771477
)
14781478

14791479
# Remove instrumentation
1480-
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
1480+
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
14811481

14821482
# Now, let's say we optimize the code and make changes.
14831483
new_fto_code = """import math
@@ -1543,7 +1543,7 @@ def risk_adjusted_return(return_val, weight):
15431543
testing_time=0.1,
15441544
)
15451545
# Remove instrumentation
1546-
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
1546+
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
15471547
matched, diffs = compare_test_results(test_results, modified_test_results)
15481548

15491549
assert not matched
@@ -1606,7 +1606,7 @@ def calculate_portfolio_metrics(
16061606
testing_time=0.1,
16071607
)
16081608
# Remove instrumentation
1609-
FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
1609+
PythonFunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path)
16101610
matched, diffs = compare_test_results(test_results, modified_test_results_2)
16111611
# now the test should match and no diffs should be found
16121612
assert len(diffs) == 0
@@ -1671,7 +1671,7 @@ def __init__(self, x, y):
16711671
file_path=sample_code_path,
16721672
parents=[FunctionParent(name="SlotsClass", type="ClassDef")],
16731673
)
1674-
func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
1674+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=fto, test_cfg=test_config)
16751675
func_optimizer.test_files = TestFiles(
16761676
test_files=[
16771677
TestFile(

tests/test_instrument_tests.py

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
TestsInFile,
2828
TestType,
2929
)
30-
from codeflash.optimization.function_optimizer import FunctionOptimizer
30+
from codeflash.languages.python.function_optimizer import PythonFunctionOptimizer
3131
from codeflash.verification.verification_utils import TestConfig
3232

3333
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():
434434
test_framework="pytest",
435435
pytest_cmd="pytest",
436436
)
437-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
437+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
438438
test_env = os.environ.copy()
439439
test_env["CODEFLASH_TEST_ITERATION"] = "0"
440440
test_env["CODEFLASH_LOOP_INDEX"] = "1"
@@ -695,7 +695,7 @@ def test_sort_parametrized(input, expected_output):
695695
test_framework="pytest",
696696
pytest_cmd="pytest",
697697
)
698-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
698+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
699699
test_results, coverage_data = func_optimizer.run_and_parse_tests(
700700
testing_type=TestingMode.BEHAVIOR,
701701
test_env=test_env,
@@ -984,7 +984,7 @@ def test_sort_parametrized_loop(input, expected_output):
984984
test_framework="pytest",
985985
pytest_cmd="pytest",
986986
)
987-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
987+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
988988
test_results, coverage_data = func_optimizer.run_and_parse_tests(
989989
testing_type=TestingMode.BEHAVIOR,
990990
test_env=test_env,
@@ -1341,7 +1341,7 @@ def test_sort():
13411341
test_framework="pytest",
13421342
pytest_cmd="pytest",
13431343
)
1344-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
1344+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
13451345
test_results, coverage_data = func_optimizer.run_and_parse_tests(
13461346
testing_type=TestingMode.BEHAVIOR,
13471347
test_env=test_env,
@@ -1723,7 +1723,7 @@ def test_sort(self):
17231723
test_framework="unittest",
17241724
pytest_cmd="pytest",
17251725
)
1726-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
1726+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
17271727
test_results, coverage_data = func_optimizer.run_and_parse_tests(
17281728
testing_type=TestingMode.BEHAVIOR,
17291729
test_env=test_env,
@@ -1973,7 +1973,7 @@ def test_sort(self, input, expected_output):
19731973
test_framework="unittest",
19741974
pytest_cmd="pytest",
19751975
)
1976-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
1976+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
19771977
test_results, coverage_data = func_optimizer.run_and_parse_tests(
19781978
testing_type=TestingMode.BEHAVIOR,
19791979
test_env=test_env,
@@ -2229,7 +2229,7 @@ def test_sort(self):
22292229
test_framework="unittest",
22302230
pytest_cmd="pytest",
22312231
)
2232-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
2232+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
22332233
test_results, coverage_data = func_optimizer.run_and_parse_tests(
22342234
test_env=test_env,
22352235
testing_type=TestingMode.BEHAVIOR,
@@ -2481,7 +2481,7 @@ def test_sort(self, input, expected_output):
24812481
test_framework="unittest",
24822482
pytest_cmd="pytest",
24832483
)
2484-
func_optimizer = FunctionOptimizer(function_to_optimize=f, test_cfg=test_config)
2484+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=f, test_cfg=test_config)
24852485
test_results, coverage_data = func_optimizer.run_and_parse_tests(
24862486
testing_type=TestingMode.BEHAVIOR,
24872487
test_env=test_env,
@@ -3144,7 +3144,7 @@ def test_sleepfunc_sequence_short(n, expected_total_sleep_time):
31443144
test_framework="pytest",
31453145
pytest_cmd="pytest",
31463146
)
3147-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
3147+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
31483148
test_files = TestFiles(
31493149
test_files=[
31503150
TestFile(
@@ -3279,7 +3279,7 @@ def test_sleepfunc_sequence_short(self, n, expected_total_sleep_time):
32793279
test_framework="unittest",
32803280
pytest_cmd="pytest",
32813281
)
3282-
func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
3282+
func_optimizer = PythonFunctionOptimizer(function_to_optimize=func, test_cfg=test_config)
32833283
test_results, coverage_data = func_optimizer.run_and_parse_tests(
32843284
testing_type=TestingMode.PERFORMANCE,
32853285
test_env=test_env,

0 commit comments

Comments
 (0)