first pass

Author: KRRT7

codeflash/cli_cmds/cli.py

@@ -75,7 +75,6 @@ def parse_args() -> Namespace:
     parser.add_argument(
         "--tests-root", type=str, help="Path to the test directory of the project, where all the tests are located."
     )
-    parser.add_argument("--test-framework", choices=["pytest", "unittest"], default="pytest")
     parser.add_argument("--config-file", type=str, help="Path to the pyproject.toml with codeflash configs.")
     parser.add_argument("--replay-test", type=str, nargs="+", help="Paths to replay test to optimize functions from")
     parser.add_argument(
@@ -172,7 +171,6 @@ def process_pyproject_config(args: Namespace) -> Namespace:
         "module_root",
         "tests_root",
         "benchmarks_root",
-        "test_framework",
         "ignore_paths",
         "pytest_cmd",
         "formatter_cmds",

codeflash/cli_cmds/cmd_init.py

@@ -1233,26 +1233,8 @@ def sorter(arr: Union[List[int],List[float]]) -> Union[List[int],List[float]]:
                 arr[j + 1] = temp
     return arr
 """
-    if args.test_framework == "unittest":
-        bubble_sort_test_content = f"""import unittest
-from {os.path.basename(args.module_root)}.bubble_sort import sorter # Keep usage of os.path.basename to avoid pathlib potential incompatibility https://github.com/codeflash-ai/codeflash/pull/1066#discussion_r1801628022
-
-class TestBubbleSort(unittest.TestCase):
-    def test_sort(self):
-        input = [5, 4, 3, 2, 1, 0]
-        output = sorter(input)
-        self.assertEqual(output, [0, 1, 2, 3, 4, 5])
-
-        input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
-        output = sorter(input)
-        self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
-
-        input = list(reversed(range(100)))
-        output = sorter(input)
-        self.assertEqual(output, list(range(100)))
-"""  # noqa: PTH119
-    elif args.test_framework == "pytest":
-        bubble_sort_test_content = f"""from {Path(args.module_root).name}.bubble_sort import sorter
+    # Always use pytest for tests
+    bubble_sort_test_content = f"""from {Path(args.module_root).name}.bubble_sort import sorter
 
 def test_sort():
     input = [5, 4, 3, 2, 1, 0]

codeflash/optimization/function_optimizer.py

@@ -1053,7 +1053,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
                     call_positions=[test.position for test in tests_in_file_list],
                     function_to_optimize=self.function_to_optimize,
                     tests_project_root=self.test_cfg.tests_project_rootdir,
-                    test_framework=self.args.test_framework,
+                    test_framework="pytest",
                 )
                 if not success:
                     continue
@@ -1063,7 +1063,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
                     call_positions=[test.position for test in tests_in_file_list],
                     function_to_optimize=self.function_to_optimize,
                     tests_project_root=self.test_cfg.tests_project_rootdir,
-                    test_framework=self.args.test_framework,
+                    test_framework="pytest",
                 )
                 if not success:
                     continue
@@ -1271,7 +1271,7 @@ def setup_and_establish_baseline(
 
         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)
+            not coverage_critic(original_code_baseline.coverage_results, "pytest")
             or not quantity_of_tests_critic(original_code_baseline)
         ):
             if self.args.override_fixtures:
@@ -1593,7 +1593,7 @@ 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
-        assert (test_framework := self.args.test_framework) in {"pytest", "unittest"}  # noqa: RUF018
+        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)
@@ -1618,7 +1618,7 @@ def establish_original_code_baseline(
                     test_files=self.test_files,
                     optimization_iteration=0,
                     testing_time=total_looping_time,
-                    enable_coverage=test_framework == "pytest",
+                    enable_coverage=True,  # Always enable coverage with pytest
                     code_context=code_context,
                 )
             finally:
@@ -1632,7 +1632,7 @@ def establish_original_code_baseline(
             )
             console.rule()
             return Failure("Failed to establish a baseline for the original code - bevhavioral tests failed.")
-        if not coverage_critic(coverage_results, self.args.test_framework):
+        if not coverage_critic(coverage_results, "pytest"):
             did_pass_all_tests = all(result.did_pass for result in behavioral_results)
             if not did_pass_all_tests:
                 return Failure("Tests failed to pass for the original code.")

codeflash/optimization/optimizer.py

@@ -44,8 +44,7 @@ def __init__(self, args: Namespace) -> None:
             tests_root=args.tests_root,
             tests_project_rootdir=args.test_project_root,
             project_root_path=args.project_root,
-            test_framework=args.test_framework,
-            pytest_cmd=args.pytest_cmd,
+            pytest_cmd=args.pytest_cmd if hasattr(args, "pytest_cmd") else "pytest",
             benchmark_tests_root=args.benchmarks_root if "benchmark" in args and "benchmarks_root" in args else None,
         )
 
@@ -274,10 +273,9 @@ def run(self) -> None:
             file_to_funcs_to_optimize, num_optimizable_functions
         )
         optimizations_found: int = 0
-        if self.args.test_framework == "pytest":
-            self.test_cfg.concolic_test_root_dir = Path(
-                tempfile.mkdtemp(dir=self.args.tests_root, prefix="codeflash_concolic_")
-            )
+        self.test_cfg.concolic_test_root_dir = Path(
+            tempfile.mkdtemp(dir=self.args.tests_root, prefix="codeflash_concolic_")
+        )
         try:
             ph("cli-optimize-functions-to-optimize", {"num_functions": num_optimizable_functions})
             if num_optimizable_functions == 0:

codeflash/verification/concolic_testing.py

@@ -81,8 +81,7 @@ def generate_concolic_tests(
                 tests_root=concolic_test_suite_dir,
                 tests_project_rootdir=test_cfg.concolic_test_root_dir,
                 project_root_path=args.project_root,
-                test_framework=args.test_framework,
-                pytest_cmd=args.pytest_cmd,
+                pytest_cmd=args.pytest_cmd if hasattr(args, "pytest_cmd") else "pytest",
             )
             function_to_concolic_tests, num_discovered_concolic_tests, _ = discover_unit_tests(concolic_test_cfg)
             logger.info(

codeflash/verification/verification_utils.py

@@ -69,11 +69,14 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
 class TestConfig:
     tests_root: Path
     project_root_path: Path
-    test_framework: str
     tests_project_rootdir: Path
-    # tests_project_rootdir corresponds to pytest rootdir,
-    # or for unittest - project_root_from_module_root(args.tests_root, pyproject_file_path)
+    # tests_project_rootdir corresponds to pytest rootdir
     concolic_test_root_dir: Optional[Path] = None
     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."""
+        return "pytest"