diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index 2a0ddc310..4c167fb69 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -2,13 +2,87 @@ import ast import os +import re import shutil import site +from contextlib import contextmanager from functools import lru_cache from pathlib import Path from tempfile import TemporaryDirectory +import tomlkit + from codeflash.cli_cmds.console import logger +from codeflash.code_utils.config_parser import find_pyproject_toml + + +@contextmanager +def custom_addopts() -> None: + pyproject_file = find_pyproject_toml() + original_content = None + non_blacklist_plugin_args = "" + + try: + # Read original file + if pyproject_file.exists(): + with Path.open(pyproject_file, encoding="utf-8") as f: + original_content = f.read() + data = tomlkit.parse(original_content) + # Backup original addopts + original_addopts = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "") + # nothing to do if no addopts present + if original_addopts != "": + original_addopts = [x.strip() for x in original_addopts] + non_blacklist_plugin_args = re.sub(r"-n(?: +|=)\S+", "", " ".join(original_addopts)).split(" ") + non_blacklist_plugin_args = [x for x in non_blacklist_plugin_args if x != ""] + if non_blacklist_plugin_args != original_addopts: + data["tool"]["pytest"]["ini_options"]["addopts"] = non_blacklist_plugin_args + # Write modified file + with Path.open(pyproject_file, "w", encoding="utf-8") as f: + f.write(tomlkit.dumps(data)) + + yield + + finally: + # Restore original file + if ( + original_content + and pyproject_file.exists() + and tuple(original_addopts) not in {(), tuple(non_blacklist_plugin_args)} + ): + with Path.open(pyproject_file, "w", encoding="utf-8") as f: + f.write(original_content) + + +@contextmanager +def add_addopts_to_pyproject() -> None: + pyproject_file = find_pyproject_toml() + original_content = None + try: + # Read original file + if pyproject_file.exists(): + with Path.open(pyproject_file, encoding="utf-8") as f: + original_content = f.read() + data = tomlkit.parse(original_content) + data["tool"]["pytest"] = {} + data["tool"]["pytest"]["ini_options"] = {} + data["tool"]["pytest"]["ini_options"]["addopts"] = [ + "-n=auto", + "-n", + "1", + "-n 1", + "-n 1", + "-n auto", + ] + with Path.open(pyproject_file, "w", encoding="utf-8") as f: + f.write(tomlkit.dumps(data)) + + yield + + finally: + # Restore original file + with Path.open(pyproject_file, "w", encoding="utf-8") as f: + f.write(original_content) def encoded_tokens_len(s: str) -> int: diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index 258092850..77a357225 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -17,7 +17,7 @@ from pydantic.dataclasses import dataclass from codeflash.cli_cmds.console import console, logger, test_files_progress_bar -from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path +from codeflash.code_utils.code_utils import custom_addopts, get_run_tmp_file, module_name_from_file_path from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_cache_db from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile, TestType @@ -150,19 +150,20 @@ def discover_tests_pytest( project_root = cfg.project_root_path tmp_pickle_path = get_run_tmp_file("collected_tests.pkl") - result = subprocess.run( - [ - SAFE_SYS_EXECUTABLE, - Path(__file__).parent / "pytest_new_process_discovery.py", - str(project_root), - str(tests_root), - str(tmp_pickle_path), - ], - cwd=project_root, - check=False, - capture_output=True, - text=True, - ) + with custom_addopts(): + result = subprocess.run( + [ + SAFE_SYS_EXECUTABLE, + Path(__file__).parent / "pytest_new_process_discovery.py", + str(project_root), + str(tests_root), + str(tmp_pickle_path), + ], + cwd=project_root, + check=False, + capture_output=True, + text=True, + ) try: with tmp_pickle_path.open(mode="rb") as f: exitcode, tests, pytest_rootdir = pickle.load(f) diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index 58834d1c4..f0d1c01a5 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING from codeflash.cli_cmds.console import logger -from codeflash.code_utils.code_utils import get_run_tmp_file +from codeflash.code_utils.code_utils import custom_addopts, get_run_tmp_file from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME from codeflash.code_utils.coverage_utils import prepare_coverage_files @@ -23,8 +23,9 @@ def execute_test_subprocess( cmd_list: list[str], cwd: Path, env: dict[str, str] | None, timeout: int = 600 ) -> subprocess.CompletedProcess: """Execute a subprocess with the given command list, working directory, environment variables, and timeout.""" - logger.debug(f"executing test run with command: {' '.join(cmd_list)}") - return subprocess.run(cmd_list, capture_output=True, cwd=cwd, env=env, text=True, timeout=timeout, check=False) + with custom_addopts(): + logger.debug(f"executing test run with command: {' '.join(cmd_list)}") + return subprocess.run(cmd_list, capture_output=True, cwd=cwd, env=env, text=True, timeout=timeout, check=False) def run_behavioral_tests( @@ -97,6 +98,7 @@ def run_behavioral_tests( coverage_cmd.extend(shlex.split(pytest_cmd, posix=IS_POSIX)[1:]) blocklist_args = [f"-p no:{plugin}" for plugin in BEHAVIORAL_BLOCKLISTED_PLUGINS if plugin != "cov"] + results = execute_test_subprocess( coverage_cmd + common_pytest_args + blocklist_args + result_args + test_files, cwd=cwd, @@ -252,7 +254,6 @@ def run_benchmarking_tests( pytest_test_env = test_env.copy() pytest_test_env["PYTEST_PLUGINS"] = "codeflash.verification.pytest_plugin" blocklist_args = [f"-p no:{plugin}" for plugin in BENCHMARKING_BLOCKLISTED_PLUGINS] - results = execute_test_subprocess( pytest_cmd_list + pytest_args + blocklist_args + result_args + test_files, cwd=cwd, @@ -278,7 +279,6 @@ def run_unittest_tests( log_level = ["-v"] if verbose else [] files = [str(file) for file in test_file_paths] output_file = ["--output-file", str(result_file_path)] - results = execute_test_subprocess( unittest_cmd_list + log_level + files + output_file, cwd=cwd, env=test_env, timeout=600 ) diff --git a/tests/scripts/end_to_end_test_topological_sort.py b/tests/scripts/end_to_end_test_topological_sort.py index ae561bc9d..2ff2bfada 100644 --- a/tests/scripts/end_to_end_test_topological_sort.py +++ b/tests/scripts/end_to_end_test_topological_sort.py @@ -1,23 +1,29 @@ import os import pathlib +import tomlkit +from codeflash.code_utils.code_utils import add_addopts_to_pyproject from end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command, run_with_retries def run_test(expected_improvement_pct: int) -> bool: - config = TestConfig( - file_path="topological_sort.py", - function_name="Graph.topologicalSort", - test_framework="pytest", - min_improvement_x=0.05, - coverage_expectations=[ - CoverageExpectation( - function_name="Graph.topologicalSort", expected_coverage=100.0, expected_lines=[24, 25, 26, 27, 28, 29] - ) - ], - ) - cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve() - return run_codeflash_command(cwd, config, expected_improvement_pct) + with add_addopts_to_pyproject(): + config = TestConfig( + file_path="topological_sort.py", + function_name="Graph.topologicalSort", + test_framework="pytest", + min_improvement_x=0.05, + coverage_expectations=[ + CoverageExpectation( + function_name="Graph.topologicalSort", + expected_coverage=100.0, + expected_lines=[24, 25, 26, 27, 28, 29], + ) + ], + ) + cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve() + return_var = run_codeflash_command(cwd, config, expected_improvement_pct) + return return_var if __name__ == "__main__": diff --git a/tests/scripts/end_to_end_test_tracer_replay.py b/tests/scripts/end_to_end_test_tracer_replay.py index 2af49999d..a2506fd99 100644 --- a/tests/scripts/end_to_end_test_tracer_replay.py +++ b/tests/scripts/end_to_end_test_tracer_replay.py @@ -10,7 +10,7 @@ def run_test(expected_improvement_pct: int) -> bool: min_improvement_x=0.1, expected_unit_tests=1, coverage_expectations=[ - CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[5, 6, 7, 8, 10, 13]), + CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[5, 6, 7, 8, 10, 13]) ], ) cwd = ( @@ -18,5 +18,6 @@ def run_test(expected_improvement_pct: int) -> bool: ).resolve() return run_codeflash_command(cwd, config, expected_improvement_pct) + if __name__ == "__main__": exit(run_with_retries(run_test, int(os.getenv("EXPECTED_IMPROVEMENT_PCT", 10)))) diff --git a/tests/scripts/end_to_end_test_utilities.py b/tests/scripts/end_to_end_test_utilities.py index d050f50e9..3645930b0 100644 --- a/tests/scripts/end_to_end_test_utilities.py +++ b/tests/scripts/end_to_end_test_utilities.py @@ -117,7 +117,9 @@ def run_codeflash_command( return validated -def build_command(cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path, benchmarks_root:pathlib.Path|None = None) -> list[str]: +def build_command( + cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path, benchmarks_root: pathlib.Path | None = None +) -> list[str]: python_path = "../../../codeflash/main.py" if "code_directories" in str(cwd) else "../codeflash/main.py" base_command = ["python", python_path, "--file", config.file_path, "--no-pr"] @@ -251,4 +253,4 @@ def run_with_retries(test_func, *args, **kwargs) -> bool: logging.error("Test failed after all retries") return 1 - return 1 \ No newline at end of file + return 1