Merge branch 'main' into ashraf/cf-894-vsc-init-no-way-to-create-tests-dir-when-no-test-dir-present

Author: mashraf-222

.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/api/aiservice.py

@@ -153,7 +153,6 @@ def optimize_python_code(  # noqa: D417
 
         if response.status_code == 200:
             optimizations_json = response.json()["optimizations"]
-            logger.info(f"!lsp|Generated {len(optimizations_json)} candidate optimizations.")
             console.rule()
             end_time = time.perf_counter()
             logger.debug(f"!lsp|Generating possible optimizations took {end_time - start_time:.2f} seconds.")

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
@@ -60,7 +59,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
@@ -71,7 +69,6 @@ class CLISetupInfo:
 class VsCodeSetupInfo:
     module_root: str
     tests_root: str
-    test_framework: str
     formatter: Union[str, list[str]]
 
 
@@ -257,7 +254,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:
@@ -293,9 +289,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}"
@@ -480,43 +473,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
@@ -613,60 +569,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)
@@ -1085,7 +994,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
@@ -1350,26 +1258,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

@@ -225,7 +225,6 @@ def get_config_value(key: str, default: str = "") -> str:
     setup_info = VsCodeSetupInfo(
         module_root=module_root,
         tests_root=tests_root,
-        test_framework=get_config_value("test_framework", "pytest"),
         formatter=get_formatter_cmds(get_config_value("formatter_cmds", "disabled")),
     )
 
@@ -241,7 +240,6 @@ def get_config_value(key: str, default: str = "") -> str:
 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()
 
@@ -276,7 +274,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/lsp/features/perform_optimization.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import concurrent.futures
 import contextlib
+import contextvars
 import os
 from typing import TYPE_CHECKING
 
-from codeflash.cli_cmds.console import code_print
+from codeflash.cli_cmds.console import code_print, logger
 from codeflash.code_utils.git_worktree_utils import create_diff_patch_from_worktree
 from codeflash.either import is_successful
 
@@ -44,24 +46,48 @@ def sync_perform_optimization(server: CodeflashLanguageServer, cancel_event: thr
         function_optimizer.function_to_tests = function_to_tests
 
     abort_if_cancelled(cancel_event)
-    test_setup_result = function_optimizer.generate_and_instrument_tests(
-        code_context, should_run_experiment=should_run_experiment
-    )
+
+    ctx_tests = contextvars.copy_context()
+    ctx_opts = contextvars.copy_context()
+
+    def run_generate_tests():  # noqa: ANN202
+        return function_optimizer.generate_and_instrument_tests(code_context)
+
+    def run_generate_optimizations():  # noqa: ANN202
+        return function_optimizer.generate_optimizations(
+            read_writable_code=code_context.read_writable_code,
+            read_only_context_code=code_context.read_only_context_code,
+            run_experiment=should_run_experiment,
+        )
+
+    future_tests = function_optimizer.executor.submit(ctx_tests.run, run_generate_tests)
+    future_optimizations = function_optimizer.executor.submit(ctx_opts.run, run_generate_optimizations)
+
+    logger.info(f"loading|Generating new tests and optimizations for function '{params.functionName}'...")
+    concurrent.futures.wait([future_tests, future_optimizations])
+
+    test_setup_result = future_tests.result()
+    optimization_result = future_optimizations.result()
+
     abort_if_cancelled(cancel_event)
     if not is_successful(test_setup_result):
         return {"functionName": params.functionName, "status": "error", "message": test_setup_result.failure()}
+    if not is_successful(optimization_result):
+        return {"functionName": params.functionName, "status": "error", "message": optimization_result.failure()}
+
     (
         generated_tests,
         function_to_concolic_tests,
         concolic_test_str,
-        optimizations_set,
         generated_test_paths,
         generated_perf_test_paths,
         instrumented_unittests_created_for_function,
         original_conftest_content,
-        function_references,
     ) = test_setup_result.unwrap()
 
+    optimizations_set, function_references = optimization_result.unwrap()
+
+    logger.info(f"Generated '{len(optimizations_set.control)}' candidate optimizations.")
     baseline_setup_result = function_optimizer.setup_and_establish_baseline(
         code_context=code_context,
         original_helper_code=original_helper_code,