test: cover real Handoff object branch in visualization (#3100)

Author: Quratulain-bilal

tests/test_visualization.py

@@ -3,7 +3,7 @@
 import graphviz  # type: ignore
 import pytest
 
-from agents import Agent
+from agents import Agent, handoff
 from agents.extensions.visualization import (
     draw_graph,
     get_all_edges,
@@ -210,3 +210,35 @@ def test_draw_graph_with_real_agent_with_handoffs():
     assert '"ParentAgent" -> "__end__"' not in graph.source
     # Child has no handoffs, so should connect to __end__
     assert '"ChildAgent" -> "__end__"' in graph.source
+
+
+def test_draw_graph_with_real_handoff_object():
+    """Test draw_graph with a real Handoff object (not just Agent) in handoffs.
+
+    Exercises the ``isinstance(handoff, Handoff)`` branches in get_all_nodes /
+    get_all_edges (rather than the ``isinstance(handoff, Agent)`` branches),
+    using the public ``handoff()`` factory rather than ``Mock(spec=Handoff)``.
+    """
+    child_agent = Agent(name="ChildAgent", instructions="Child instructions")
+    real_handoff = handoff(child_agent)
+    assert isinstance(real_handoff, Handoff)
+
+    parent_agent = Agent(
+        name="ParentAgent",
+        instructions="Parent instructions",
+        handoffs=[real_handoff],
+    )
+
+    graph = draw_graph(parent_agent)
+
+    assert isinstance(graph, graphviz.Source)
+    assert '"ParentAgent"' in graph.source
+    # Node uses agent_name from the Handoff object
+    assert (
+        '"ChildAgent" [label="ChildAgent", shape=box, style=filled, style=rounded, '
+        "fillcolor=lightyellow, width=1.5, height=0.8];" in graph.source
+    )
+    # Edge points from parent to handoff agent_name
+    assert '"ParentAgent" -> "ChildAgent";' in graph.source
+    # Parent has handoffs, so should NOT connect directly to __end__
+    assert '"ParentAgent" -> "__end__"' not in graph.source