single pass for num_discovered_tests

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -317,16 +317,15 @@ def discover_unit_tests(
     if file_to_funcs_to_optimize:
         functions_to_optimize = [func for funcs_list in file_to_funcs_to_optimize.values() for func in funcs_list]
 
-    function_to_tests = strategy(cfg, discover_only_these_tests, functions_to_optimize)
-    num_discovered_tests = sum([len(value) for value in function_to_tests.values()])
+    function_to_tests, num_discovered_tests = strategy(cfg, discover_only_these_tests, functions_to_optimize)
     return function_to_tests, num_discovered_tests
 
 
 def discover_tests_pytest(
     cfg: TestConfig,
     discover_only_these_tests: list[Path] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[str, set[FunctionCalledInTest]]:
+) -> tuple[dict[str, set[FunctionCalledInTest]], int]:
     tests_root = cfg.tests_root
     project_root = cfg.project_root_path
 
@@ -402,7 +401,7 @@ def discover_tests_unittest(
     cfg: TestConfig,
     discover_only_these_tests: list[str] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[str, set[FunctionCalledInTest]]:
+) -> tuple[dict[str, set[FunctionCalledInTest]], int]:
     tests_root: Path = cfg.tests_root
     loader: unittest.TestLoader = unittest.TestLoader()
     tests: unittest.TestSuite = loader.discover(str(tests_root))
@@ -469,7 +468,7 @@ def process_test_files(
     file_to_test_map: dict[Path, list[TestsInFile]],
     cfg: TestConfig,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[str, set[FunctionCalledInTest]]:
+) -> tuple[dict[str, set[FunctionCalledInTest]], int]:
     import jedi
 
     project_root_path = cfg.project_root_path
@@ -492,6 +491,7 @@ def process_test_files(
         logger.debug(f"Import analysis results: {len(import_results)} files analyzed")
 
     function_to_test_map = defaultdict(set)
+    num_discovered_tests = 0
     jedi_project = jedi.Project(path=project_root_path)
 
     with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
@@ -639,7 +639,8 @@ def process_test_files(
                                 position=CodePosition(line_no=name.line, col_no=name.column),
                             )
                         )
+                        num_discovered_tests += 1
 
             progress.advance(task_id)
 
-    return dict(function_to_test_map)
+    return dict(function_to_test_map), num_discovered_tests