Merge branch 'main' into saga4/micro_fix_12

Author: aseembits93

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:

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)

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
     )

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__":

tests/scripts/end_to_end_test_tracer_replay.py

@@ -10,13 +10,14 @@ 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 = (
         pathlib.Path(__file__).parent.parent.parent / "code_to_optimize" / "code_directories" / "simple_tracer_e2e"
     ).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))))

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
+    return 1