|
6 | 6 | import pytest |
7 | 7 | from codeflash.context.unused_definition_remover import detect_unused_helper_functions |
8 | 8 | from codeflash.discovery.functions_to_optimize import FunctionToOptimize |
| 9 | +from codeflash.models.models import get_code_block_splitter |
9 | 10 | from codeflash.optimization.function_optimizer import FunctionOptimizer |
10 | 11 | from codeflash.verification.verification_utils import TestConfig |
11 | 12 |
|
@@ -55,7 +56,8 @@ def test_detect_unused_helper_functions(temp_project): |
55 | 56 | temp_dir, main_file, test_cfg = temp_project |
56 | 57 |
|
57 | 58 | # Optimized version that only calls one helper |
58 | | - optimized_code = """ |
| 59 | + optimized_code = f""" |
| 60 | +{get_code_block_splitter("main.py")} |
59 | 61 | def entrypoint_function(n): |
60 | 62 | \"\"\"Optimized function that only calls one helper.\"\"\" |
61 | 63 | result1 = helper_function_1(n) |
@@ -99,7 +101,8 @@ def helper_function_2(x): |
99 | 101 |
|
100 | 102 | # Also test the complete replace_function_and_helpers_with_optimized_code workflow |
101 | 103 | # First modify the optimized code to include a MODIFIED unused helper |
102 | | - optimized_code_with_modified_helper = """ |
| 104 | + optimized_code_with_modified_helper = f""" |
| 105 | +{get_code_block_splitter("main.py")} |
103 | 106 | def entrypoint_function(n): |
104 | 107 | \"\"\"Optimized function that only calls one helper.\"\"\" |
105 | 108 | result1 = helper_function_1(n) |
@@ -159,7 +162,8 @@ def test_revert_unused_helper_functions(temp_project): |
159 | 162 | temp_dir, main_file, test_cfg = temp_project |
160 | 163 |
|
161 | 164 | # Optimized version that only calls one helper and modifies the unused one |
162 | | - optimized_code = """ |
| 165 | + optimized_code = f""" |
| 166 | +{get_code_block_splitter("main.py") } |
163 | 167 | def entrypoint_function(n): |
164 | 168 | \"\"\"Optimized function that only calls one helper.\"\"\" |
165 | 169 | result1 = helper_function_1(n) |
@@ -221,7 +225,8 @@ def test_no_unused_helpers_no_revert(temp_project): |
221 | 225 | temp_dir, main_file, test_cfg = temp_project |
222 | 226 |
|
223 | 227 | # Optimized version that still calls both helpers |
224 | | - optimized_code = """ |
| 228 | + optimized_code = f""" |
| 229 | +{get_code_block_splitter("main.py")} |
225 | 230 | def entrypoint_function(n): |
226 | 231 | \"\"\"Optimized function that still calls both helpers.\"\"\" |
227 | 232 | result1 = helper_function_1(n) |
@@ -303,13 +308,16 @@ def helper_function_2(x): |
303 | 308 | """) |
304 | 309 |
|
305 | 310 | # Optimized version that only calls one helper |
306 | | - optimized_code = """ |
| 311 | + optimized_code = f""" |
| 312 | +{get_code_block_splitter("main.py")} |
307 | 313 | from helpers import helper_function_1 |
308 | 314 |
|
309 | 315 | def entrypoint_function(n): |
310 | 316 | \"\"\"Optimized function that only calls one helper.\"\"\" |
311 | 317 | result1 = helper_function_1(n) |
312 | 318 | return result1 + n * 3 # Inlined helper_function_2 |
| 319 | +
|
| 320 | +{get_code_block_splitter("helpers.py")} |
313 | 321 | """ |
314 | 322 |
|
315 | 323 | # Create test config |
@@ -384,7 +392,6 @@ def helper_function_2(x): |
384 | 392 |
|
385 | 393 | # Apply optimization and test reversion |
386 | 394 | optimizer.replace_function_and_helpers_with_optimized_code(code_context, optimized_code, original_helper_code) |
387 | | - |
388 | 395 | # Check main file content |
389 | 396 | main_content = main_file.read_text() |
390 | 397 | assert "result1 + n * 3" in main_content, "Entrypoint function should be optimized" |
@@ -478,7 +485,8 @@ def helper_method_2(self, x): |
478 | 485 | """) |
479 | 486 |
|
480 | 487 | # Optimized version that only calls one helper method |
481 | | - optimized_code = """ |
| 488 | + optimized_code = f""" |
| 489 | +{get_code_block_splitter("main.py") } |
482 | 490 | class Calculator: |
483 | 491 | def entrypoint_method(self, n): |
484 | 492 | \"\"\"Optimized method that only calls one helper.\"\"\" |
@@ -537,7 +545,8 @@ def helper_method_2(self, x): |
537 | 545 |
|
538 | 546 | # Also test the complete replace_function_and_helpers_with_optimized_code workflow |
539 | 547 | # Update optimized code to include a MODIFIED unused helper |
540 | | - optimized_code_with_modified_helper = """ |
| 548 | + optimized_code_with_modified_helper = f""" |
| 549 | +{get_code_block_splitter("main.py")} |
541 | 550 | class Calculator: |
542 | 551 | def entrypoint_method(self, n): |
543 | 552 | \"\"\"Optimized method that only calls one helper.\"\"\" |
@@ -619,7 +628,8 @@ def process_data(self, n): |
619 | 628 | """) |
620 | 629 |
|
621 | 630 | # Optimized version that only calls one external helper |
622 | | - optimized_code = """ |
| 631 | + optimized_code = f""" |
| 632 | +{get_code_block_splitter("main.py") } |
623 | 633 | def external_helper_1(x): |
624 | 634 | \"\"\"External helper function.\"\"\" |
625 | 635 | return x * 2 |
@@ -678,7 +688,8 @@ def process_data(self, n): |
678 | 688 |
|
679 | 689 | # Also test the complete replace_function_and_helpers_with_optimized_code workflow |
680 | 690 | # Update optimized code to include a MODIFIED unused helper |
681 | | - optimized_code_with_modified_helper = """ |
| 691 | + optimized_code_with_modified_helper = f""" |
| 692 | +{get_code_block_splitter("main.py")} |
682 | 693 | def external_helper_1(x): |
683 | 694 | \"\"\"External helper function.\"\"\" |
684 | 695 | return x * 2 |
@@ -716,7 +727,8 @@ def process_data(self, n): |
716 | 727 |
|
717 | 728 | # Also test the complete replace_function_and_helpers_with_optimized_code workflow |
718 | 729 | # Update optimized code to include a MODIFIED unused helper |
719 | | - optimized_code_with_modified_helper = """ |
| 730 | + optimized_code_with_modified_helper = f""" |
| 731 | +{get_code_block_splitter("main.py")} |
720 | 732 | def external_helper_1(x): |
721 | 733 | \"\"\"External helper function.\"\"\" |
722 | 734 | return x * 2 |
@@ -786,7 +798,8 @@ def local_helper(self, x): |
786 | 798 | """) |
787 | 799 |
|
788 | 800 | # Optimized version that inlines one helper |
789 | | - optimized_code = """ |
| 801 | + optimized_code = f""" |
| 802 | +{get_code_block_splitter("main.py") } |
790 | 803 | def global_helper_1(x): |
791 | 804 | return x * 2 |
792 | 805 |
|
@@ -954,7 +967,8 @@ def clean_data(x): |
954 | 967 | """) |
955 | 968 |
|
956 | 969 | # Optimized version that only uses some functions |
957 | | - optimized_code = """ |
| 970 | + optimized_code = f""" |
| 971 | +{get_code_block_splitter("main.py") } |
958 | 972 | import utils |
959 | 973 | from math_helpers import add |
960 | 974 |
|
@@ -1115,7 +1129,8 @@ def divide_numbers(x, y): |
1115 | 1129 | """) |
1116 | 1130 |
|
1117 | 1131 | # Optimized version that only uses add_numbers |
1118 | | - optimized_code = """ |
| 1132 | + optimized_code = f""" |
| 1133 | +{get_code_block_splitter("main.py") } |
1119 | 1134 | import calculator |
1120 | 1135 |
|
1121 | 1136 | def entrypoint_function(n): |
@@ -1317,7 +1332,8 @@ def calculate_class(cls, n): |
1317 | 1332 | """) |
1318 | 1333 |
|
1319 | 1334 | # Optimized static method that inlines one utility |
1320 | | - optimized_static_code = """ |
| 1335 | + optimized_static_code = f""" |
| 1336 | +{get_code_block_splitter("main.py")} |
1321 | 1337 | def utility_function_1(x): |
1322 | 1338 | return x * 2 |
1323 | 1339 |
|
@@ -1384,7 +1400,8 @@ def calculate_class(cls, n): |
1384 | 1400 |
|
1385 | 1401 | # Also test the complete replace_function_and_helpers_with_optimized_code workflow |
1386 | 1402 | # Update optimized code to include a MODIFIED unused helper |
1387 | | - optimized_static_code_with_modified_helper = """ |
| 1403 | + optimized_static_code_with_modified_helper = f""" |
| 1404 | +{get_code_block_splitter("main.py")} |
1388 | 1405 | def utility_function_1(x): |
1389 | 1406 | return x * 2 |
1390 | 1407 |
|
|
0 commit comments