move progress bar

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -13,7 +13,7 @@
 from pydantic.dataclasses import dataclass
 from pytest import ExitCode
 
-from codeflash.cli_cmds.console import console, logger, progress_bar
+from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
 from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile
@@ -49,8 +49,8 @@ def discover_unit_tests(
     if not strategy:
         error_message = f"Unsupported test framework: {cfg.test_framework}"
         raise ValueError(error_message)
-    with progress_bar("Discovering unit tests…", transient=True):
-        return strategy(cfg, discover_only_these_tests)
+
+    return strategy(cfg, discover_only_these_tests)
 
 
 def discover_tests_pytest(

codeflash/optimization/optimizer.py

@@ -8,7 +8,7 @@
 from typing import TYPE_CHECKING
 
 from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient
-from codeflash.cli_cmds.console import console, logger
+from codeflash.cli_cmds.console import console, logger, progress_bar
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_replacer import normalize_code, normalize_node
 from codeflash.code_utils.code_utils import get_run_tmp_file
@@ -95,9 +95,8 @@ def run(self) -> None:
                 return
 
             console.rule()
-            logger.info(f"Discovering existing unit tests in {self.test_cfg.tests_root}…")
-            console.rule()
-            function_to_tests: dict[str, list[FunctionCalledInTest]] = discover_unit_tests(self.test_cfg)
+            with progress_bar(f"Discovering existing unit tests in {self.test_cfg.tests_root}…", transient=True):
+                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()])
             console.rule()
             logger.info(f"Discovered {num_discovered_tests} existing unit tests in {self.test_cfg.tests_root}")