Update discover_unit_tests.py

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -182,11 +182,13 @@ def visit_Call(self, node: ast.Call) -> None:
             # __import__("module_name")
             if isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
                 self.imported_modules.add(node.args[0].value)
-        elif (isinstance(node.func, ast.Attribute)
-              and isinstance(node.func.value, ast.Name)
-              and node.func.value.id == "importlib"
-              and node.func.attr == "import_module"
-              and node.args):
+        elif (
+            isinstance(node.func, ast.Attribute)
+            and isinstance(node.func.value, ast.Name)
+            and node.func.value.id == "importlib"
+            and node.func.attr == "import_module"
+            and node.args
+        ):
             # importlib.import_module("module_name")
             if isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
                 self.imported_modules.add(node.args[0].value)
@@ -263,8 +265,7 @@ def analyze_imports_in_test_file(test_file_path: Path, target_functions: set[str
 
 
 def filter_test_files_by_imports(
-    file_to_test_map: dict[Path, list[TestsInFile]],
-    target_functions: set[str]
+    file_to_test_map: dict[Path, list[TestsInFile]], target_functions: set[str]
 ) -> tuple[dict[Path, list[TestsInFile]], dict[Path, set[str]]]:
     """Filter test files based on import analysis to reduce Jedi processing.
 
@@ -297,7 +298,9 @@ 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
+    cfg: TestConfig,
+    discover_only_these_tests: list[Path] | None = None,
+    functions_to_optimize: 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)
@@ -309,7 +312,9 @@ def discover_unit_tests(
 
 
 def discover_tests_pytest(
-    cfg: TestConfig, discover_only_these_tests: list[Path] | None = None, functions_to_optimize: list[FunctionToOptimize] | None = None
+    cfg: TestConfig,
+    discover_only_these_tests: list[Path] | None = None,
+    functions_to_optimize: list[FunctionToOptimize] | None = None,
 ) -> dict[Path, list[FunctionCalledInTest]]:
     tests_root = cfg.tests_root
     project_root = cfg.project_root_path
@@ -383,7 +388,9 @@ def discover_tests_pytest(
 
 
 def discover_tests_unittest(
-    cfg: TestConfig, discover_only_these_tests: list[str] | None = None, functions_to_optimize: list[FunctionToOptimize] | None = None
+    cfg: TestConfig,
+    discover_only_these_tests: list[str] | None = None,
+    functions_to_optimize: list[FunctionToOptimize] | None = None,
 ) -> dict[Path, list[FunctionCalledInTest]]:
     tests_root: Path = cfg.tests_root
     loader: unittest.TestLoader = unittest.TestLoader()
@@ -448,7 +455,9 @@ def discover_parameters_unittest(function_name: str) -> tuple[bool, str, str | N
 
 
 def process_test_files(
-    file_to_test_map: dict[Path, list[TestsInFile]], cfg: TestConfig, functions_to_optimize: list[FunctionToOptimize] | None = None
+    file_to_test_map: dict[Path, list[TestsInFile]],
+    cfg: TestConfig,
+    functions_to_optimize: list[FunctionToOptimize] | None = None,
 ) -> dict[str, list[FunctionCalledInTest]]:
     import jedi
 
@@ -466,7 +475,7 @@ def process_test_files(
             # Also add qualified name without module
             if func.parents:
                 target_function_names.add(f"{func.parents[0].name}.{func.function_name}")
-        
+
         logger.debug(f"Target functions for import filtering: {target_function_names}")
         file_to_test_map, import_results = filter_test_files_by_imports(file_to_test_map, target_function_names)
         logger.debug(f"Import analysis results: {len(import_results)} files analyzed")