start testing now

Author: aseembits93

code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py

@@ -3,29 +3,10 @@
 from code_to_optimize.bubble_sort import sorter
 
 
-class DummyBenchmark:
-    def __call__(self, func, *args, **kwargs):
-        # Mimic calling benchmark(func, *args, **kwargs)
-        return func(*args, **kwargs)
-
-    def __getattr__(self, name):
-        # Mimic benchmark attributes like .pedantic, .extra_info etc.
-        def dummy(*args, **kwargs):
-            return None
-
-        return dummy
-
-
-@pytest.fixture
-def benchmark():
-    return DummyBenchmark()
-
-
 def test_sort(benchmark):
     result = benchmark(sorter, list(reversed(range(500))))
     assert result == list(range(500))
 
-
 # This should not be picked up as a benchmark test
 def test_sort2():
     result = sorter(list(reversed(range(500))))

codeflash/code_utils/code_replacer.py

@@ -76,26 +76,27 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
 
 
 @contextlib.contextmanager
-def disable_autouse(test_path: Path) -> None:
+def disable_autouse(test_path: Path) -> str:
     file_content = test_path.read_text(encoding="utf-8")
-    try:
-        module = cst.parse_module(file_content)
-        disable_autouse_fixture = AutouseFixtureModifier()
-        modified_module = module.visit(disable_autouse_fixture)
-        test_path.write_text(modified_module.code, encoding="utf-8")
-    finally:
-        test_path.write_text(file_content, encoding="utf-8")
+    module = cst.parse_module(file_content)
+    disable_autouse_fixture = AutouseFixtureModifier()
+    modified_module = module.visit(disable_autouse_fixture)
+    test_path.write_text(modified_module.code, encoding="utf-8")
+    return file_content
 
 
-def modify_autouse_fixture(test_paths: list[Path]) -> None:
+def modify_autouse_fixture(test_paths: list[Path]) -> dict[Path, list[str]]:
     # find fixutre definition in conftetst.py (the one closest to the test)
     # get fixtures present in override-fixtures in pyproject.toml
     # add if marker closest return
+    file_content_map = {}
     conftest_files = find_conftest_files(test_paths)
     for cf_file in conftest_files:
         # iterate over all functions in the file
         # if function has autouse fixture, modify function to bypass with custom marker
-        disable_autouse(cf_file)
+        original_content = disable_autouse(cf_file)
+        file_content_map[cf_file] = original_content
+    return file_content_map
 
 
 # # reuse line profiler utils to add decorator and import to test fns

codeflash/code_utils/code_utils.py

@@ -208,3 +208,8 @@ def cleanup_paths(paths: list[Path]) -> None:
                 shutil.rmtree(path, ignore_errors=True)
             else:
                 path.unlink(missing_ok=True)
+
+
+def restore_conftest(path_to_content_map: dict[Path, str]) -> None:
+    for path, file_content in path_to_content_map.items():
+        path.write_text(file_content, encoding="utf8")

codeflash/optimization/function_optimizer.py

@@ -33,6 +33,7 @@
     get_run_tmp_file,
     has_any_async_functions,
     module_name_from_file_path,
+    restore_conftest,
 )
 from codeflash.code_utils.config_consts import (
     INDIVIDUAL_TESTCASE_TIMEOUT,
@@ -212,7 +213,8 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
             for key in set(self.function_to_tests) | set(function_to_concolic_tests)
         }
         instrumented_unittests_created_for_function = self.instrument_existing_tests(function_to_all_tests)
-
+        # logger.debug("disabling all autouse fixtures associated with the test files")
+        original_conftest_content = modify_autouse_fixture(list(instrumented_unittests_created_for_function))
         # Get a dict of file_path_to_classes of fto and helpers_of_fto
         file_path_to_helper_classes = defaultdict(set)
         for function_source in code_context.helper_functions:
@@ -234,13 +236,15 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
         )
 
         if not is_successful(baseline_result):
+            restore_conftest(original_conftest_content)
             cleanup_paths(paths_to_cleanup)
             return Failure(baseline_result.failure())
 
         original_code_baseline, test_functions_to_remove = baseline_result.unwrap()
         if isinstance(original_code_baseline, OriginalCodeBaseline) and not coverage_critic(
             original_code_baseline.coverage_results, self.args.test_framework
         ):
+            restore_conftest(original_conftest_content)
             cleanup_paths(paths_to_cleanup)
             return Failure("The threshold for test coverage was not met.")
         # request for new optimizations but don't block execution, check for completion later
@@ -746,8 +750,6 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi
                 f"{concolic_coverage_test_files_count} concolic coverage test file"
                 f"{'s' if concolic_coverage_test_files_count != 1 else ''} for {func_qualname}"
             )
-        logger.debug("disabling all autouse fixtures associated with the test files")
-        modify_autouse_fixture(list(unique_instrumented_test_files))
         logger.debug("add custom marker to all tests")
         add_custom_marker_to_all_tests(list(unique_instrumented_test_files))
         return unique_instrumented_test_files