fix for E2E

Author: KRRT7

codeflash/optimization/function_optimizer.py

@@ -1591,7 +1591,6 @@ def establish_original_code_baseline(
     ) -> Result[tuple[OriginalCodeBaseline, list[str]], str]:
         line_profile_results = {"timings": {}, "unit": 0, "str_out": ""}
         # For the original function - run the tests and get the runtime, plus coverage
-        test_framework = "pytest"  # Always use pytest for all tests
         success = True
 
         test_env = self.get_test_env(codeflash_loop_index=0, codeflash_test_iteration=0, codeflash_tracer_disable=1)

codeflash/verification/verification_utils.py

@@ -75,7 +75,7 @@ class TestConfig:
     pytest_cmd: str = "pytest"
     benchmark_tests_root: Optional[Path] = None
     use_cache: bool = True
-    
+
     @property
     def test_framework(self) -> str:
         """Always returns 'pytest' as we use pytest for all tests."""

tests/scripts/end_to_end_test_benchmark_sort.py

@@ -10,7 +10,6 @@ def run_test(expected_improvement_pct: int) -> bool:
         file_path=pathlib.Path("bubble_sort.py"),
         function_name="sorter",
         benchmarks_root=cwd / "tests" / "pytest" / "benchmarks",
-        test_framework="pytest",
         min_improvement_x=0.70,
         coverage_expectations=[
             CoverageExpectation(

tests/scripts/end_to_end_test_bubblesort_pytest.py

@@ -8,7 +8,6 @@ def run_test(expected_improvement_pct: int) -> bool:
     config = TestConfig(
         file_path="bubble_sort.py",
         function_name="sorter",
-        test_framework="pytest",
         min_improvement_x=0.70,
         coverage_expectations=[
             CoverageExpectation(

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", test_framework="unittest", min_improvement_x=0.30
+        file_path="bubble_sort.py", function_name="sorter", min_improvement_x=0.30
     )
     cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve()
     return run_codeflash_command(cwd, config, expected_improvement_pct)

tests/scripts/end_to_end_test_init_optimization.py

@@ -8,7 +8,6 @@ def run_test(expected_improvement_pct: int) -> bool:
     config = TestConfig(
         file_path="remove_control_chars.py",
         function_name="CharacterRemover.remove_control_characters",
-        test_framework="pytest",
         min_improvement_x=0.1,
         coverage_expectations=[
             CoverageExpectation(

tests/scripts/end_to_end_test_topological_sort_worktree.py

@@ -8,7 +8,6 @@ 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,
         use_worktree=True,
         coverage_expectations=[

tests/scripts/end_to_end_test_utilities.py

@@ -21,7 +21,6 @@ class TestConfig:
     # Make file_path optional when trace_mode is True
     file_path: Optional[pathlib.Path] = None
     function_name: Optional[str] = None
-    test_framework: Optional[str] = None
     expected_unit_tests: Optional[int] = None
     min_improvement_x: float = 0.1
     trace_mode: bool = False
@@ -86,7 +85,7 @@ def run_codeflash_command(
 
     path_to_file = cwd / config.file_path
     file_contents = path_to_file.read_text("utf-8")
-    test_root = cwd / "tests" / (config.test_framework or "")
+    test_root = cwd / "tests" / "pytest"  # Always use pytest
 
     command = build_command(cwd, config, test_root, config.benchmarks_root if config.benchmarks_root else None)
     env = os.environ.copy()
@@ -129,10 +128,8 @@ def build_command(
 
     if config.function_name:
         base_command.extend(["--function", config.function_name])
-    if config.test_framework:
-        base_command.extend(
-            ["--test-framework", config.test_framework, "--tests-root", str(test_root), "--module-root", str(cwd)]
-        )
+    # Always use pytest - test framework is specified via tests-root and module-root
+    base_command.extend(["--tests-root", str(test_root), "--module-root", str(cwd)])
     if benchmarks_root:
         base_command.extend(["--benchmark", "--benchmarks-root", str(benchmarks_root)])
     if config.use_worktree:
@@ -190,7 +187,7 @@ def validate_stdout_in_candidate(stdout: str, expected_in_stdout: list[str]) ->
 
 
 def run_trace_test(cwd: pathlib.Path, config: TestConfig, expected_improvement_pct: int) -> bool:
-    test_root = cwd / "tests" / (config.test_framework or "")
+    test_root = cwd / "tests" / "pytest"  # Always use pytest
     clear_directory(test_root)
     command = ["uv", "run", "--no-project", "-m", "codeflash.main", "optimize", "workload.py"]
     env = os.environ.copy()