moved the num_discovered_tests calculation inside the discover_unit_tests

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -305,7 +305,7 @@ def discover_unit_tests(
     cfg: TestConfig,
     discover_only_these_tests: list[Path] | None = None,
     file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] | None = None,
-) -> dict[str, list[FunctionCalledInTest]]:
+) -> tuple[dict[str, list[FunctionCalledInTest]], int]:
     framework_strategies: dict[str, Callable] = {"pytest": discover_tests_pytest, "unittest": discover_tests_unittest}
     strategy = framework_strategies.get(cfg.test_framework, None)
     if not strategy:
@@ -317,7 +317,9 @@ 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]
 
-    return strategy(cfg, discover_only_these_tests, functions_to_optimize)
+    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()])
+    return function_to_tests, num_discovered_tests
 
 
 def discover_tests_pytest(

codeflash/discovery/functions_to_optimize.py

@@ -268,7 +268,7 @@ def find_all_functions_in_file(file_path: Path) -> dict[Path, list[FunctionToOpt
 def get_all_replay_test_functions(
     replay_test: Path, test_cfg: TestConfig, project_root_path: Path
 ) -> dict[Path, list[FunctionToOptimize]]:
-    function_tests = discover_unit_tests(test_cfg, discover_only_these_tests=[replay_test])
+    function_tests, _ = discover_unit_tests(test_cfg, discover_only_these_tests=[replay_test])
     # Get the absolute file paths for each function, excluding class name if present
     filtered_valid_functions = defaultdict(list)
     file_to_functions_map = defaultdict(list)

codeflash/optimization/optimizer.py

@@ -162,8 +162,9 @@ def run(self) -> None:
 
             console.rule()
             start_time = time.time()
-            function_to_tests = discover_unit_tests(self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize)
-            num_discovered_tests: int = sum([len(value) for value in function_to_tests.values()])
+            function_to_tests, num_discovered_tests = discover_unit_tests(
+                self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize
+            )
             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}"

codeflash/verification/concolic_testing.py

@@ -78,8 +78,7 @@ def generate_concolic_tests(
                 test_framework=args.test_framework,
                 pytest_cmd=args.pytest_cmd,
             )
-            function_to_concolic_tests = discover_unit_tests(concolic_test_cfg)
-            num_discovered_concolic_tests: int = sum([len(value) for value in function_to_concolic_tests.values()])
+            function_to_concolic_tests, num_discovered_concolic_tests = discover_unit_tests(concolic_test_cfg)
             logger.info(
                 f"Created {num_discovered_concolic_tests} "
                 f"concolic unit test case{'s' if num_discovered_concolic_tests != 1 else ''} "

tests/test_unit_test_discovery.py

@@ -20,7 +20,7 @@ def test_unit_test_discovery_pytest():
         test_framework="pytest",
         tests_project_rootdir=tests_path.parent,
     )
-    tests = discover_unit_tests(test_config)
+    tests, _ = discover_unit_tests(test_config)
     assert len(tests) > 0
 
 
@@ -33,7 +33,7 @@ def test_benchmark_test_discovery_pytest():
         test_framework="pytest",
         tests_project_rootdir=tests_path.parent,
     )
-    tests = discover_unit_tests(test_config)
+    tests, _ = discover_unit_tests(test_config)
     assert len(tests) == 1 # Should not discover benchmark tests
 
 
@@ -47,7 +47,7 @@ def test_unit_test_discovery_unittest():
         tests_project_rootdir=project_path.parent,
     )
     os.chdir(project_path)
-    tests = discover_unit_tests(test_config)
+    tests, _ = discover_unit_tests(test_config)
     # assert len(tests) > 0
     # Unittest discovery within a pytest environment does not work
 
@@ -85,7 +85,7 @@ def sorter(arr):
         )
 
         # Discover tests
-        tests = discover_unit_tests(test_config)
+        tests, _ = discover_unit_tests(test_config)
         assert len(tests) == 1
         assert 'bubble_sort.sorter' in tests
         assert len(tests['bubble_sort.sorter']) == 2
@@ -124,7 +124,7 @@ def test_discover_tests_pytest_with_temp_dir_root():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the dummy test file is discovered
         assert len(discovered_tests) == 1
@@ -197,7 +197,7 @@ def test_discover_tests_pytest_with_multi_level_dirs():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test files at all levels are discovered
         assert len(discovered_tests) == 3
@@ -287,7 +287,7 @@ def test_discover_tests_pytest_dirs():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test files at all levels are discovered
         assert len(discovered_tests) == 4
@@ -333,7 +333,7 @@ def test_discover_tests_pytest_with_class():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test class and method are discovered
         assert len(discovered_tests) == 1
@@ -371,7 +371,7 @@ def test_discover_tests_pytest_with_double_nested_directories():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test class and method are discovered
         assert len(discovered_tests) == 1
@@ -421,7 +421,7 @@ def test_discover_tests_with_code_in_dir_and_test_in_subdir():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test file is discovered and associated with the code file
         assert len(discovered_tests) == 1
@@ -460,7 +460,7 @@ def test_discover_tests_pytest_with_nested_class():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test for the nested class method is discovered
         assert len(discovered_tests) == 1
@@ -500,7 +500,7 @@ def test_discover_tests_pytest_separate_moduledir():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test for the nested class method is discovered
         assert len(discovered_tests) == 1
@@ -542,7 +542,7 @@ def test_add(self):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the unittest was discovered
         assert len(discovered_tests) == 1
@@ -609,7 +609,7 @@ def test_add(self):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the unittest was discovered
         assert len(discovered_tests) == 2
@@ -654,7 +654,7 @@ def _test_add(self):  # Private test method should not be discovered
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify no tests were discovered
         assert len(discovered_tests) == 0
@@ -706,7 +706,7 @@ def test_add_with_parameters(self):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the unittest was discovered
         assert len(discovered_tests) == 1
@@ -788,7 +788,7 @@ def test_add_mixed(self, name, a, b, expected):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the basic structure
         assert len(discovered_tests) == 2  # Should have tests for both add and multiply
@@ -1069,7 +1069,7 @@ def test_other():
         )
 
         # Test without filtering
-        all_tests = discover_unit_tests(test_config)
+        all_tests, _ = discover_unit_tests(test_config)
         assert len(all_tests) == 2  # Should find both functions
 
         # Test with filtering - create mock FunctionToOptimize objects
@@ -1079,7 +1079,7 @@ def test_other():
         mock_function.function_name = "target_function"
         mock_function.parents = []  # No parent classes
 
-        filtered_tests = discover_unit_tests(test_config, file_to_funcs_to_optimize={code_file: [mock_function]})
+        filtered_tests, _ = discover_unit_tests(test_config, file_to_funcs_to_optimize={code_file: [mock_function]})
         assert len(filtered_tests) >= 1
         assert "mycode.target_function" in filtered_tests
 
@@ -1212,7 +1212,7 @@ def test_unrelated():
         )
 
         # Test without filtering
-        all_tests = discover_unit_tests(test_config)
+        all_tests, _ = discover_unit_tests(test_config)
         assert len(all_tests) == 2  # Should find both functions
 
         # Test with filtering - create mock FunctionToOptimize objects
@@ -1222,7 +1222,7 @@ def test_unrelated():
         mock_function.function_name = "target_function"
         mock_function.parents = []  # No parent classes
 
-        filtered_tests = discover_unit_tests(test_config, file_to_funcs_to_optimize={target_file: [mock_function]})
+        filtered_tests, _ = discover_unit_tests(test_config, file_to_funcs_to_optimize={target_file: [mock_function]})
         # Should filter out the unrelated test since it imports from a different module
         assert len(filtered_tests) == 1
         assert "target_module.target_function" in filtered_tests