diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index a6e28aaaa..271eb06bc 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -80,6 +80,9 @@ def parse_args() -> Namespace: parser.add_argument( "--no-pr", action="store_true", help="Do not create a PR for the optimization, only update the code locally." ) + parser.add_argument( + "--no-gen-tests", action="store_true", help="Do not generate tests, use only existing tests for optimization." + ) parser.add_argument("--staging-review", action="store_true", help="Upload optimizations to staging for review") parser.add_argument( "--verify-setup", diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 416bdc8df..50399839d 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -1295,47 +1295,51 @@ def generate_tests( n_tests = N_TESTS_TO_GENERATE_EFFECTIVE assert len(generated_test_paths) == n_tests - # Submit test generation tasks - future_tests = self.submit_test_generation_tasks( - self.executor, - testgen_context.markdown, - [definition.fully_qualified_name for definition in helper_functions], - generated_test_paths, - generated_perf_test_paths, - ) + if not self.args.no_gen_tests: + # Submit test generation tasks + future_tests = self.submit_test_generation_tasks( + self.executor, + testgen_context.markdown, + [definition.fully_qualified_name for definition in helper_functions], + generated_test_paths, + generated_perf_test_paths, + ) future_concolic_tests = self.executor.submit( generate_concolic_tests, self.test_cfg, self.args, self.function_to_optimize, self.function_to_optimize_ast ) - # Wait for test futures to complete - concurrent.futures.wait([*future_tests, future_concolic_tests]) - + if not self.args.no_gen_tests: + # Wait for test futures to complete + concurrent.futures.wait([*future_tests, future_concolic_tests]) + else: + concurrent.futures.wait([future_concolic_tests]) # Process test generation results tests: list[GeneratedTests] = [] - for future in future_tests: - res = future.result() - if res: - ( - generated_test_source, - instrumented_behavior_test_source, - instrumented_perf_test_source, - test_behavior_path, - test_perf_path, - ) = res - tests.append( - GeneratedTests( - generated_original_test_source=generated_test_source, - instrumented_behavior_test_source=instrumented_behavior_test_source, - instrumented_perf_test_source=instrumented_perf_test_source, - behavior_file_path=test_behavior_path, - perf_file_path=test_perf_path, + if not self.args.no_gen_tests: + for future in future_tests: + res = future.result() + if res: + ( + generated_test_source, + instrumented_behavior_test_source, + instrumented_perf_test_source, + test_behavior_path, + test_perf_path, + ) = res + tests.append( + GeneratedTests( + generated_original_test_source=generated_test_source, + instrumented_behavior_test_source=instrumented_behavior_test_source, + instrumented_perf_test_source=instrumented_perf_test_source, + behavior_file_path=test_behavior_path, + perf_file_path=test_perf_path, + ) ) - ) - if not tests: - logger.warning(f"Failed to generate and instrument tests for {self.function_to_optimize.function_name}") - return Failure(f"/!\\ NO TESTS GENERATED for {self.function_to_optimize.function_name}") + if not tests: + logger.warning(f"Failed to generate and instrument tests for {self.function_to_optimize.function_name}") + return Failure(f"/!\\ NO TESTS GENERATED for {self.function_to_optimize.function_name}") function_to_concolic_tests, concolic_test_str = future_concolic_tests.result() count_tests = len(tests) diff --git a/tests/scripts/end_to_end_test_bubblesort_unittest.py b/tests/scripts/end_to_end_test_bubblesort_unittest.py index c7f0107b1..ad717253f 100644 --- a/tests/scripts/end_to_end_test_bubblesort_unittest.py +++ b/tests/scripts/end_to_end_test_bubblesort_unittest.py @@ -6,7 +6,7 @@ def run_test(expected_improvement_pct: int) -> bool: config = TestConfig( - file_path="bubble_sort.py", function_name="sorter", min_improvement_x=0.30 + file_path="bubble_sort.py", function_name="sorter", min_improvement_x=0.30, no_gen_tests=True ) cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve() return run_codeflash_command(cwd, config, expected_improvement_pct) diff --git a/tests/scripts/end_to_end_test_utilities.py b/tests/scripts/end_to_end_test_utilities.py index 777bf16ba..d9c95d7e3 100644 --- a/tests/scripts/end_to_end_test_utilities.py +++ b/tests/scripts/end_to_end_test_utilities.py @@ -36,6 +36,7 @@ class TestConfig: coverage_expectations: list[CoverageExpectation] = field(default_factory=list) benchmarks_root: Optional[pathlib.Path] = None use_worktree: bool = False + no_gen_tests: bool = False def clear_directory(directory_path: str | pathlib.Path) -> None: @@ -84,6 +85,11 @@ def validate_coverage(stdout: str, expectations: list[CoverageExpectation]) -> b return True +def validate_no_gen_tests(stdout: str) -> bool: + if "Generated '0' tests for" not in stdout: + logging.error("Tests generated even when flag was on") + return False + return True def run_codeflash_command( cwd: pathlib.Path, config: TestConfig, expected_improvement_pct: int, expected_in_stdout: list[str] = None @@ -156,6 +162,8 @@ def build_command( base_command.extend(["--benchmark", "--benchmarks-root", str(benchmarks_root)]) if config.use_worktree: base_command.append("--worktree") + if config.no_gen_tests: + base_command.append("--no-gen-tests") return base_command @@ -214,6 +222,9 @@ def validate_output(stdout: str, return_code: int, expected_improvement_pct: int if config.coverage_expectations: validate_coverage(stdout, config.coverage_expectations) + if config.no_gen_tests: + validate_no_gen_tests(stdout) + logging.info(f"Success: Performance improvement is {improvement_pct}%") return True