ready

Author: Codeflash Bot

codeflash/discovery/discover_unit_tests.py

@@ -296,6 +296,17 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                     self.found_qualified_name = target_func
                     return
 
+            # Check if any target function is a method of the imported class/module
+            # e.g., importing Graph and looking for Graph.topologicalSort
+            # TODO will pick up all tests which have the same class name (could be coming from a different file)
+            for target_func in fnames:
+                if "." in target_func:
+                    class_name, method_name = target_func.split(".", 1)
+                    if aname == class_name:
+                        self.found_any_target_function = True
+                        self.found_qualified_name = target_func
+                        return
+
     def visit_Attribute(self, node: ast.Attribute) -> None:
         """Handle attribute access like module.function_name."""
         if self.found_any_target_function:

tests/scripts/end_to_end_test_topological_sort_worktree.py

@@ -18,6 +18,7 @@ def run_test(expected_improvement_pct: int) -> bool:
                 expected_lines=[25, 26, 27, 28, 29, 30, 31],
             )
         ],
+        expected_unit_tests=1,
     )
     cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve()
     return_var = run_codeflash_command(cwd, config, expected_improvement_pct)