parallelize test discovery

Author: KRRT7

codeflash/code_utils/checkpoint.py

@@ -8,7 +8,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional
 
-import click
+from rich.prompt import Confirm
 
 if TYPE_CHECKING:
     import argparse
@@ -137,7 +137,7 @@ def ask_should_use_checkpoint_get_functions(args: argparse.Namespace) -> Optiona
     previous_checkpoint_functions = None
     if args.all and (sys.platform == "linux" or sys.platform == "darwin") and Path("/tmp").is_dir():  # noqa: S108 #TODO: use the temp dir from codeutils-compat.py
         previous_checkpoint_functions = get_all_historical_functions(args.module_root, Path("/tmp"))  # noqa: S108
-        if previous_checkpoint_functions and click.confirm(
+        if previous_checkpoint_functions and Confirm.ask(
             "Previous Checkpoint detected from an incomplete optimization run, shall I continue the optimization from that point?",
             default=True,
         ):

codeflash/optimization/optimizer.py

@@ -162,8 +162,8 @@ def run(self) -> None:
 
             console.rule()
             start_time = time.time()
-            function_to_tests: dict[str, list[FunctionCalledInTest]] = discover_unit_tests(self.test_cfg)
-            num_discovered_tests: int = sum([len(value) for value in function_to_tests.values()])
+            function_to_tests: dict[str, list[FunctionCalledInTest]]
+            function_to_tests, num_discovered_tests = discover_unit_tests(self.test_cfg)
             console.rule()
             logger.info(
                 f"Discovered {num_discovered_tests} existing unit tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"