fix: Add integ tests

Author: Murat Kaan Meral

tests_integ/test_multiagent_graph.py

@@ -1,3 +1,5 @@
+from typing import Any, AsyncIterator
+
 import pytest
 
 from strands import Agent, tool
@@ -9,6 +11,7 @@
     BeforeModelCallEvent,
     MessageAddedEvent,
 )
+from strands.multiagent.base import MultiAgentBase, MultiAgentResult, NodeResult, Status
 from strands.multiagent.graph import GraphBuilder
 from strands.types.content import ContentBlock
 from tests.fixtures.mock_hook_provider import MockHookProvider
@@ -218,3 +221,182 @@ async def test_graph_execution_with_image(image_analysis_agent, summary_agent, y
 
     assert hook_provider.extract_for(image_analysis_agent).event_types_received == expected_hook_events
     assert hook_provider.extract_for(summary_agent).event_types_received == expected_hook_events
+
+
+class CustomStreamingNode(MultiAgentBase):
+    """Custom node that wraps an agent and adds custom streaming events."""
+
+    def __init__(self, agent: Agent, name: str):
+        self.agent = agent
+        self.name = name
+
+    async def invoke_async(
+        self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any
+    ) -> MultiAgentResult:
+        result = await self.agent.invoke_async(task, **kwargs)
+        node_result = NodeResult(result=result, status=Status.COMPLETED)
+        return MultiAgentResult(status=Status.COMPLETED, results={self.name: node_result})
+
+    async def stream_async(
+        self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any
+    ) -> AsyncIterator[dict[str, Any]]:
+        yield {"custom_event": "start", "node": self.name}
+        result = await self.agent.invoke_async(task, **kwargs)
+        yield {"custom_event": "agent_complete", "node": self.name}
+        node_result = NodeResult(result=result, status=Status.COMPLETED)
+        yield {"result": MultiAgentResult(status=Status.COMPLETED, results={self.name: node_result})}
+
+
+@pytest.mark.asyncio
+async def test_graph_streaming_with_agents():
+    """Test that Graph properly streams events from agent nodes."""
+    math_agent = Agent(
+        name="math",
+        model="us.amazon.nova-pro-v1:0",
+        system_prompt="You are a math assistant.",
+        tools=[calculate_sum],
+    )
+    summary_agent = Agent(
+        name="summary",
+        model="us.amazon.nova-lite-v1:0",
+        system_prompt="You are a summary assistant.",
+    )
+
+    builder = GraphBuilder()
+    builder.add_node(math_agent, "math")
+    builder.add_node(summary_agent, "summary")
+    builder.add_edge("math", "summary")
+    builder.set_entry_point("math")
+    graph = builder.build()
+
+    # Collect events
+    events = []
+    async for event in graph.stream_async("Calculate 5 + 3 and summarize the result"):
+        events.append(event)
+
+    # Count event categories
+    node_start_events = [e for e in events if e.get("multi_agent_node_start")]
+    node_stream_events = [e for e in events if e.get("multi_agent_node_stream")]
+    node_complete_events = [e for e in events if e.get("multi_agent_node_complete")]
+    result_events = [e for e in events if "result" in e and "multi_agent_node_start" not in e]
+
+    # Verify we got multiple events of each type
+    assert len(node_start_events) >= 2, f"Expected at least 2 node_start events, got {len(node_start_events)}"
+    assert len(node_stream_events) > 10, f"Expected many node_stream events, got {len(node_stream_events)}"
+    assert len(node_complete_events) >= 2, f"Expected at least 2 node_complete events, got {len(node_complete_events)}"
+    assert len(result_events) >= 1, f"Expected at least 1 result event, got {len(result_events)}"
+
+    # Verify we have events for both nodes
+    math_events = [e for e in events if e.get("node_id") == "math"]
+    summary_events = [e for e in events if e.get("node_id") == "summary"]
+    assert len(math_events) > 0, "Expected events from math node"
+    assert len(summary_events) > 0, "Expected events from summary node"
+
+
+@pytest.mark.asyncio
+async def test_graph_streaming_with_custom_node():
+    """Test that Graph properly streams events from custom MultiAgentBase nodes."""
+    math_agent = Agent(
+        name="math",
+        model="us.amazon.nova-pro-v1:0",
+        system_prompt="You are a math assistant.",
+        tools=[calculate_sum],
+    )
+    summary_agent = Agent(
+        name="summary",
+        model="us.amazon.nova-lite-v1:0",
+        system_prompt="You are a summary assistant.",
+    )
+
+    # Create a custom node
+    custom_node = CustomStreamingNode(summary_agent, "custom_summary")
+
+    builder = GraphBuilder()
+    builder.add_node(math_agent, "math")
+    builder.add_node(custom_node, "custom_summary")
+    builder.add_edge("math", "custom_summary")
+    builder.set_entry_point("math")
+    graph = builder.build()
+
+    # Collect events
+    events = []
+    async for event in graph.stream_async("Calculate 5 + 3 and summarize the result"):
+        events.append(event)
+
+    # Count event categories
+    node_start_events = [e for e in events if e.get("multi_agent_node_start")]
+    node_stream_events = [e for e in events if e.get("multi_agent_node_stream")]
+    custom_events = [e for e in events if e.get("custom_event")]
+    result_events = [e for e in events if "result" in e and "multi_agent_node_start" not in e]
+
+    # Verify we got multiple events of each type
+    assert len(node_start_events) >= 2, f"Expected at least 2 node_start events, got {len(node_start_events)}"
+    assert len(node_stream_events) > 5, f"Expected many node_stream events, got {len(node_stream_events)}"
+    assert len(custom_events) >= 2, f"Expected at least 2 custom events (start, complete), got {len(custom_events)}"
+    assert len(result_events) >= 1, f"Expected at least 1 result event, got {len(result_events)}"
+
+    # Verify custom events are properly structured
+    custom_start = [e for e in custom_events if e.get("custom_event") == "start"]
+    custom_complete = [e for e in custom_events if e.get("custom_event") == "agent_complete"]
+
+    assert len(custom_start) >= 1, "Expected at least 1 custom start event"
+    assert len(custom_complete) >= 1, "Expected at least 1 custom complete event"
+
+
+@pytest.mark.asyncio
+async def test_nested_graph_streaming():
+    """Test that nested graphs properly propagate streaming events."""
+    math_agent = Agent(
+        name="math",
+        model="us.amazon.nova-pro-v1:0",
+        system_prompt="You are a math assistant.",
+        tools=[calculate_sum],
+    )
+    analysis_agent = Agent(
+        name="analysis",
+        model="us.amazon.nova-lite-v1:0",
+        system_prompt="You are an analysis assistant.",
+    )
+
+    # Create nested graph
+    nested_builder = GraphBuilder()
+    nested_builder.add_node(math_agent, "calculator")
+    nested_builder.add_node(analysis_agent, "analyzer")
+    nested_builder.add_edge("calculator", "analyzer")
+    nested_builder.set_entry_point("calculator")
+    nested_graph = nested_builder.build()
+
+    # Create outer graph with nested graph
+    summary_agent = Agent(
+        name="summary",
+        model="us.amazon.nova-lite-v1:0",
+        system_prompt="You are a summary assistant.",
+    )
+
+    outer_builder = GraphBuilder()
+    outer_builder.add_node(nested_graph, "computation")
+    outer_builder.add_node(summary_agent, "summary")
+    outer_builder.add_edge("computation", "summary")
+    outer_builder.set_entry_point("computation")
+    outer_graph = outer_builder.build()
+
+    # Collect events
+    events = []
+    async for event in outer_graph.stream_async("Calculate 7 + 8 and provide a summary"):
+        events.append(event)
+
+    # Count event categories
+    node_start_events = [e for e in events if e.get("multi_agent_node_start")]
+    node_stream_events = [e for e in events if e.get("multi_agent_node_stream")]
+    result_events = [e for e in events if "result" in e and "multi_agent_node_start" not in e]
+
+    # Verify we got multiple events
+    assert len(node_start_events) >= 2, f"Expected at least 2 node_start events, got {len(node_start_events)}"
+    assert len(node_stream_events) > 10, f"Expected many node_stream events, got {len(node_stream_events)}"
+    assert len(result_events) >= 1, f"Expected at least 1 result event, got {len(result_events)}"
+
+    # Verify we have events from nested nodes
+    computation_events = [e for e in events if e.get("node_id") == "computation"]
+    summary_events = [e for e in events if e.get("node_id") == "summary"]
+    assert len(computation_events) > 0, "Expected events from computation (nested graph) node"
+    assert len(summary_events) > 0, "Expected events from summary node"

tests_integ/test_multiagent_swarm.py

@@ -134,3 +134,41 @@ async def test_swarm_execution_with_image(researcher_agent, analyst_agent, write
 
     # Verify agent history - at least one agent should have been used
     assert len(result.node_history) > 0
+
+
+@pytest.mark.asyncio
+async def test_swarm_streaming():
+    """Test that Swarm properly streams events during execution."""
+    researcher = Agent(
+        name="researcher",
+        model="us.amazon.nova-pro-v1:0",
+        system_prompt="You are a researcher. When you need calculations, hand off to the analyst.",
+    )
+    analyst = Agent(
+        name="analyst",
+        model="us.amazon.nova-pro-v1:0",
+        system_prompt="You are an analyst. Use tools to perform calculations.",
+        tools=[calculate],
+    )
+
+    swarm = Swarm([researcher, analyst])
+
+    # Collect events
+    events = []
+    async for event in swarm.stream_async("Calculate 10 + 5 and explain the result"):
+        events.append(event)
+
+    # Count event categories
+    node_start_events = [e for e in events if e.get("multi_agent_node_start")]
+    node_stream_events = [e for e in events if e.get("multi_agent_node_stream")]
+    result_events = [e for e in events if "result" in e and "multi_agent_node_start" not in e]
+
+    # Verify we got multiple events of each type
+    assert len(node_start_events) >= 1, f"Expected at least 1 node_start event, got {len(node_start_events)}"
+    assert len(node_stream_events) > 10, f"Expected many node_stream events, got {len(node_stream_events)}"
+    assert len(result_events) >= 1, f"Expected at least 1 result event, got {len(result_events)}"
+
+    # Verify we have events from at least one agent
+    researcher_events = [e for e in events if e.get("node_id") == "researcher"]
+    analyst_events = [e for e in events if e.get("node_id") == "analyst"]
+    assert len(researcher_events) > 0 or len(analyst_events) > 0, "Expected events from at least one agent"