fix replay test detection

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -150,11 +150,13 @@ def __init__(self, function_names_to_find: set[str]) -> None:
         self.imported_modules: set[str] = set()
         self.has_dynamic_imports: bool = False
         self.wildcard_modules: set[str] = set()
+        # Track aliases: alias_name -> original_name
+        self.alias_mapping: dict[str, str] = {}
 
         # Precompute function_names for prefix search
         # For prefix match, store mapping from prefix-root to candidates for O(1) matching
         self._exact_names = function_names_to_find
-        self._prefix_roots = {}
+        self._prefix_roots: dict[str, list[str]] = {}
         for name in function_names_to_find:
             if "." in name:
                 root = name.split(".", 1)[0]
@@ -223,6 +225,12 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                 return
 
             # Fast prefix match: only for relevant roots
+            for target_func in fnames:
+                if target_func.startswith(f"{aname}."):
+                    self.found_any_target_function = True
+                    self.found_qualified_name = target_func
+                    return
+
             prefix = qname + "."
             # Only bother if one of the targets startswith the prefix-root
             candidates = proots.get(qname, ())
@@ -247,6 +255,15 @@ def visit_Attribute(self, node: ast.Attribute) -> None:
             self.found_qualified_name = node.attr
             return
 
+        if isinstance(node.value, ast.Name) and node.value.id in self.imported_modules:
+            for target_func in self.function_names_to_find:
+                if "." in target_func:
+                    class_name, method_name = target_func.rsplit(".", 1)
+                    if node.attr == method_name:
+                        self.found_any_target_function = True
+                        self.found_qualified_name = target_func
+                        return
+
         # Check if this is accessing a target function through a dynamically imported module
         # Only if we've detected dynamic imports are being used
         if self.has_dynamic_imports and node.attr in self.function_names_to_find: