move things around

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -304,22 +304,29 @@ def filter_test_files_by_imports(
 def discover_unit_tests(
     cfg: TestConfig,
     discover_only_these_tests: list[Path] | None = None,
-    functions_to_optimize: list[FunctionToOptimize] | None = None,
+    file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] | None = None,
 ) -> dict[str, list[FunctionCalledInTest]]:
     framework_strategies: dict[str, Callable] = {"pytest": discover_tests_pytest, "unittest": discover_tests_unittest}
     strategy = framework_strategies.get(cfg.test_framework, None)
     if not strategy:
         error_message = f"Unsupported test framework: {cfg.test_framework}"
         raise ValueError(error_message)
 
+    # Extract all functions to optimize for import filtering
+    functions_to_optimize = None
+    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)
 
 
 def discover_tests_pytest(
     cfg: TestConfig,
     discover_only_these_tests: list[Path] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[Path, list[FunctionCalledInTest]]:
+) -> dict[str, list[FunctionCalledInTest]]:
     tests_root = cfg.tests_root
     project_root = cfg.project_root_path
 
@@ -395,7 +402,7 @@ def discover_tests_unittest(
     cfg: TestConfig,
     discover_only_these_tests: list[str] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[Path, list[FunctionCalledInTest]]:
+) -> dict[str, list[FunctionCalledInTest]]:
     tests_root: Path = cfg.tests_root
     loader: unittest.TestLoader = unittest.TestLoader()
     tests: unittest.TestSuite = loader.discover(str(tests_root))

codeflash/optimization/optimizer.py

@@ -162,12 +162,8 @@ def run(self) -> None:
 
             console.rule()
             start_time = time.time()
-            # Extract all functions to optimize for import filtering
-            all_functions_to_optimize = [
-                func for funcs_list in file_to_funcs_to_optimize.values() for func in funcs_list
-            ]
             function_to_tests: dict[str, list[FunctionCalledInTest]] = discover_unit_tests(
-                self.test_cfg, functions_to_optimize=all_functions_to_optimize
+                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()])
             console.rule()