remove test_framework from pyproject.toml (#955)

* follow up

* remove requirement

* Delete uv.lock

* refresh uv-lock

* first pass

* cleanup test_framework here

* cleanup

* code_review

* cleanup tests

* fix for E2E

* fix tests dir missing

* one more cleanup

* cancel-in-progress

* Revert "cancel-in-progress"

This reverts commit f4bb907.

* not needed here

* lower threshold and cleanup comments

* debug

* temp

* debug

Revert "debug"

This reverts commit fc36551.

fix(discover): Fix pytest discovery for futurehouse structure

Revert "fix(discover): Fix pytest discovery for futurehouse structure"

This reverts commit 40c48882b7413f5876af0e2e08d8f17a65bab091.

Reapply "debug"

This reverts commit c8297e5.

Revert "not needed here"

This reverts commit dd2c5cd.

Revert "lower threshold and cleanup comments"

This reverts commit 0e2f57e.

Reapply "lower threshold and cleanup comments"

This reverts commit e3b24f4a2967551eca8a19f96bf6647b23acdbbc.

Reapply "not needed here"

This reverts commit aec32103c931ff6d57dfa0d012113c2cec5d37a7.

Revert "Reapply "debug""

This reverts commit 77ab9f34f858a17fb29764c544769a0eb72ce7f0.

Reapply "fix(discover): Fix pytest discovery for futurehouse structure"

This reverts commit 506b94ab4fe17a7c8e0d458253812758cced3f22.

feat(futurehouse): Make futurehouse structure pytest compatible

* Revert "debug"

This reverts commit 271c5a3.

* Revert "temp"

This reverts commit b363acd.

* Revert "debug"

This reverts commit ac29b6b.

* just for now

Author: KRRT7

.github/workflows/e2e-bubblesort-unittest.yaml

@@ -65,7 +65,6 @@ jobs:
       - name: Install dependencies (CLI)
         run: |
           uv sync
-          uv add timeout_decorator
 
       - name: Run Codeflash to optimize code
         id: optimize_code

.github/workflows/e2e-futurehouse-structure.yaml

@@ -24,7 +24,7 @@ jobs:
       COLUMNS: 110
       MAX_RETRIES: 3
       RETRY_DELAY: 5
-      EXPECTED_IMPROVEMENT_PCT: 10
+      EXPECTED_IMPROVEMENT_PCT: 5
       CODEFLASH_END_TO_END: 1
     steps:
       - name: 🛎️ Checkout

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

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import ast
 import os
 import re
 import subprocess
@@ -59,7 +58,6 @@ class CLISetupInfo:
     module_root: str
     tests_root: str
     benchmarks_root: Union[str, None]
-    test_framework: str
     ignore_paths: list[str]
     formatter: Union[str, list[str]]
     git_remote: str
@@ -70,7 +68,6 @@ class CLISetupInfo:
 class VsCodeSetupInfo:
     module_root: str
     tests_root: str
-    test_framework: str
     formatter: Union[str, list[str]]
 
 
@@ -256,7 +253,6 @@ def __init__(self) -> None:
 class CommonSections(Enum):
     module_root = "module_root"
     tests_root = "tests_root"
-    test_framework = "test_framework"
     formatter_cmds = "formatter_cmds"
 
     def get_toml_key(self) -> str:
@@ -292,9 +288,6 @@ def get_suggestions(section: str) -> tuple[list[str], Optional[str]]:
     if section == CommonSections.tests_root:
         default = "tests" if "tests" in valid_subdirs else None
         return valid_subdirs, default
-    if section == CommonSections.test_framework:
-        auto_detected = detect_test_framework_from_config_files(Path.cwd())
-        return ["pytest", "unittest"], auto_detected
     if section == CommonSections.formatter_cmds:
         return ["disabled", "ruff", "black"], "disabled"
     msg = f"Unknown section: {section}"
@@ -455,43 +448,6 @@ def collect_setup_info() -> CLISetupInfo:
 
     ph("cli-tests-root-provided")
 
-    test_framework_choices, detected_framework = get_suggestions(CommonSections.test_framework)
-    autodetected_test_framework = detected_framework or detect_test_framework_from_test_files(tests_root)
-
-    framework_message = "⚗️ Let's configure your test framework.\n\n"
-    if autodetected_test_framework:
-        framework_message += f"I detected that you're using {autodetected_test_framework}. "
-    framework_message += "Please confirm or select a different one."
-
-    framework_panel = Panel(Text(framework_message, style="blue"), title="⚗️ Test Framework", border_style="bright_blue")
-    console.print(framework_panel)
-    console.print()
-
-    framework_choices = []
-    # add icons based on the detected framework
-    for choice in test_framework_choices:
-        if choice == "pytest":
-            framework_choices.append(("🧪 pytest", "pytest"))
-        elif choice == "unittest":
-            framework_choices.append(("🐍 unittest", "unittest"))
-
-    framework_questions = [
-        inquirer.List(
-            "test_framework",
-            message="Which test framework do you use?",
-            choices=framework_choices,
-            default=autodetected_test_framework or "pytest",
-            carousel=True,
-        )
-    ]
-
-    framework_answers = inquirer.prompt(framework_questions, theme=CodeflashTheme())
-    if not framework_answers:
-        apologize_and_exit()
-    test_framework = framework_answers["test_framework"]
-
-    ph("cli-test-framework-provided", {"test_framework": test_framework})
-
     benchmarks_root = None
 
     # TODO: Implement other benchmark framework options
@@ -588,60 +544,13 @@ def collect_setup_info() -> CLISetupInfo:
         module_root=str(module_root),
         tests_root=str(tests_root),
         benchmarks_root=str(benchmarks_root) if benchmarks_root else None,
-        test_framework=cast("str", test_framework),
         ignore_paths=ignore_paths,
         formatter=cast("str", formatter),
         git_remote=str(git_remote),
         enable_telemetry=enable_telemetry,
     )
 
 
-def detect_test_framework_from_config_files(curdir: Path) -> Optional[str]:
-    test_framework = None
-    pytest_files = ["pytest.ini", "pyproject.toml", "tox.ini", "setup.cfg"]
-    pytest_config_patterns = {
-        "pytest.ini": "[pytest]",
-        "pyproject.toml": "[tool.pytest.ini_options]",
-        "tox.ini": "[pytest]",
-        "setup.cfg": "[tool:pytest]",
-    }
-    for pytest_file in pytest_files:
-        file_path = curdir / pytest_file
-        if file_path.exists():
-            with file_path.open(encoding="utf8") as file:
-                contents = file.read()
-                if pytest_config_patterns[pytest_file] in contents:
-                    test_framework = "pytest"
-                    break
-        test_framework = "pytest"
-    return test_framework
-
-
-def detect_test_framework_from_test_files(tests_root: Path) -> Optional[str]:
-    test_framework = None
-    # Check if any python files contain a class that inherits from unittest.TestCase
-    for filename in tests_root.iterdir():
-        if filename.suffix == ".py":
-            with filename.open(encoding="utf8") as file:
-                contents = file.read()
-                try:
-                    node = ast.parse(contents)
-                except SyntaxError:
-                    continue
-                if any(
-                    isinstance(item, ast.ClassDef)
-                    and any(
-                        (isinstance(base, ast.Attribute) and base.attr == "TestCase")
-                        or (isinstance(base, ast.Name) and base.id == "TestCase")
-                        for base in item.bases
-                    )
-                    for item in node.body
-                ):
-                    test_framework = "unittest"
-                    break
-    return test_framework
-
-
 def check_for_toml_or_setup_file() -> str | None:
     click.echo()
     click.echo("Checking for pyproject.toml or setup.py…\r", nl=False)
@@ -1060,7 +969,6 @@ def configure_pyproject_toml(
     else:
         codeflash_section["module-root"] = setup_info.module_root
         codeflash_section["tests-root"] = setup_info.tests_root
-        codeflash_section["test-framework"] = setup_info.test_framework
         codeflash_section["ignore-paths"] = setup_info.ignore_paths
         if not setup_info.enable_telemetry:
             codeflash_section["disable-telemetry"] = not setup_info.enable_telemetry
@@ -1325,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/code_utils/config_parser.py

@@ -149,10 +149,6 @@ def parse_config_file(
         else:
             config[key] = []
 
-    if config.get("test-framework"):
-        assert config["test-framework"] in {"pytest", "unittest"}, (
-            "In pyproject.toml, Codeflash only supports the 'test-framework' as pytest and unittest."
-        )
     # see if this is happening during GitHub actions setup
     if config.get("formatter-cmds") and len(config.get("formatter-cmds")) > 0 and not override_formatter_check:
         assert config.get("formatter-cmds")[0] != "your-formatter $file", (

codeflash/code_utils/instrument_existing_tests.py

@@ -70,7 +70,6 @@ def __init__(
         self,
         function: FunctionToOptimize,
         module_path: str,
-        test_framework: str,
         call_positions: list[CodePosition],
         mode: TestingMode = TestingMode.BEHAVIOR,
     ) -> None:
@@ -79,7 +78,6 @@ def __init__(
         self.class_name = None
         self.only_function_name = function.function_name
         self.module_path = module_path
-        self.test_framework = test_framework
         self.call_positions = call_positions
         if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
             self.class_name = function.top_level_parent_name
@@ -475,7 +473,6 @@ def __init__(
         self,
         function: FunctionToOptimize,
         module_path: str,
-        test_framework: str,
         call_positions: list[CodePosition],
         mode: TestingMode = TestingMode.BEHAVIOR,
     ) -> None:
@@ -484,7 +481,6 @@ def __init__(
         self.class_name = None
         self.only_function_name = function.function_name
         self.module_path = module_path
-        self.test_framework = test_framework
         self.call_positions = call_positions
         self.did_instrument = False
         # Track function call count per test function
@@ -639,7 +635,6 @@ def inject_async_profiling_into_existing_test(
     call_positions: list[CodePosition],
     function_to_optimize: FunctionToOptimize,
     tests_project_root: Path,
-    test_framework: str,
     mode: TestingMode = TestingMode.BEHAVIOR,
 ) -> tuple[bool, str | None]:
     """Inject profiling for async function calls by setting environment variables before each call."""
@@ -657,7 +652,7 @@ def inject_async_profiling_into_existing_test(
     import_visitor.visit(tree)
     func = import_visitor.imported_as
 
-    async_instrumenter = AsyncCallInstrumenter(func, test_module_path, test_framework, call_positions, mode=mode)
+    async_instrumenter = AsyncCallInstrumenter(func, test_module_path, call_positions, mode=mode)
     tree = async_instrumenter.visit(tree)
 
     if not async_instrumenter.did_instrument:
@@ -675,12 +670,11 @@ def inject_profiling_into_existing_test(
     call_positions: list[CodePosition],
     function_to_optimize: FunctionToOptimize,
     tests_project_root: Path,
-    test_framework: str,
     mode: TestingMode = TestingMode.BEHAVIOR,
 ) -> tuple[bool, str | None]:
     if function_to_optimize.is_async:
         return inject_async_profiling_into_existing_test(
-            test_path, call_positions, function_to_optimize, tests_project_root, test_framework, mode
+            test_path, call_positions, function_to_optimize, tests_project_root, mode
         )
 
     with test_path.open(encoding="utf8") as f:
@@ -696,7 +690,7 @@ def inject_profiling_into_existing_test(
     import_visitor.visit(tree)
     func = import_visitor.imported_as
 
-    tree = InjectPerfOnly(func, test_module_path, test_framework, call_positions, mode=mode).visit(tree)
+    tree = InjectPerfOnly(func, test_module_path, call_positions, mode=mode).visit(tree)
     new_imports = [
         ast.Import(names=[ast.alias(name="time")]),
         ast.Import(names=[ast.alias(name="gc")]),

codeflash/lsp/beta.py

@@ -187,7 +187,6 @@ def write_config(params: WriteConfigParams) -> dict[str, any]:
     setup_info = VsCodeSetupInfo(
         module_root=getattr(cfg, "module_root", ""),
         tests_root=getattr(cfg, "tests_root", ""),
-        test_framework=getattr(cfg, "test_framework", "pytest"),
         formatter=get_formatter_cmds(getattr(cfg, "formatter_cmds", "disabled")),
     )
 
@@ -203,7 +202,6 @@ def write_config(params: WriteConfigParams) -> dict[str, any]:
 def get_config_suggestions(_params: any) -> dict[str, any]:
     module_root_suggestions, default_module_root = get_suggestions(CommonSections.module_root)
     tests_root_suggestions, default_tests_root = get_suggestions(CommonSections.tests_root)
-    test_framework_suggestions, default_test_framework = get_suggestions(CommonSections.test_framework)
     formatter_suggestions, default_formatter = get_suggestions(CommonSections.formatter_cmds)
     get_valid_subdirs.cache_clear()
 
@@ -238,7 +236,6 @@ def get_config_suggestions(_params: any) -> dict[str, any]:
     return {
         "module_root": {"choices": module_root_suggestions, "default": default_module_root},
         "tests_root": {"choices": tests_root_suggestions, "default": default_tests_root},
-        "test_framework": {"choices": test_framework_suggestions, "default": default_test_framework},
         "formatter_cmds": {"choices": formatter_suggestions, "default": default_formatter},
     }
 

codeflash/optimization/function_optimizer.py

@@ -1051,7 +1051,6 @@ 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,
                 )
                 if not success:
                     continue
@@ -1061,7 +1060,6 @@ 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,
                 )
                 if not success:
                     continue
@@ -1275,7 +1273,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)
             or not quantity_of_tests_critic(original_code_baseline)
         ):
             if self.args.override_fixtures:
@@ -1597,7 +1595,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
-        assert (test_framework := self.args.test_framework) in {"pytest", "unittest"}  # noqa: RUF018
         success = True
 
         test_env = self.get_test_env(codeflash_loop_index=0, codeflash_test_iteration=0, codeflash_tracer_disable=1)
@@ -1622,7 +1619,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,
                     code_context=code_context,
                 )
             finally:
@@ -1636,7 +1633,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):
             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

@@ -45,8 +45,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") and args.pytest_cmd else "pytest",
             benchmark_tests_root=args.benchmarks_root if "benchmark" in args and "benchmarks_root" in args else None,
         )
 
@@ -285,10 +284,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: