fixing tests

Author: RoxyFarhad

examples/tutorials/10_async/00_base/000_hello_acp/tests/test_agent.py

@@ -108,25 +108,41 @@ async def test_send_event_and_stream(self, client: AsyncAgentex, agent_id: str):
         """Test sending an event and streaming the response."""
         task_response = await client.agents.create_task(agent_id, params=ParamsCreateTaskRequest(name=uuid.uuid1().hex))
         task = task_response.result
+
         assert task is not None
+        task_creation_found = False
+        # Poll for the initial task creation message
+        async for message in poll_messages(
+            client=client,
+            task_id=task.id,
+            timeout=30,
+            sleep_interval=1.0,
+        ):
+            assert isinstance(message, TaskMessage)
+            if message.content and message.content.type == "text" and message.content.author == "agent":
+                assert "Hello! I've received your task" in message.content.content
+                task_creation_found = True
+                break
+
+        assert task_creation_found, "Task creation message not found in poll"
 
         user_message = "Hello, this is a test message!"
+        stream_timeout = 10
 
         # Collect events from stream
         all_events = []
 
         # Flags to track what we've received
-        task_creation_found = False
         user_echo_found = False
         agent_response_found = False
 
-        async def collect_stream_events():
-            nonlocal task_creation_found, user_echo_found, agent_response_found
+        async def collect_stream_events() -> None:
+            nonlocal user_echo_found, agent_response_found
 
             async for event in stream_agent_response(
                 client=client,
                 task_id=task.id,
-                timeout=30,
+                timeout=stream_timeout,
             ):
                 all_events.append(event)
                 # Check events as they arrive
@@ -136,21 +152,19 @@ async def collect_stream_events():
                     if content.get("content") is None:
                         continue  # Skip empty content
                     if content.get("type") == "text" and content.get("author") == "agent":
-                        # Check for initial task creation message
-                        if "Hello! I've received your task" in content.get("content", ""):
-                            task_creation_found = True
                         # Check for agent response to user message
-                        elif "Hello! I've received your message" in content.get("content", ""):
+                        if "Hello! I've received your message" in content.get("content", ""):
                             # Agent response should come after user echo
                             assert user_echo_found, "Agent response arrived before user message echo (incorrect order)"
                             agent_response_found = True
                     elif content.get("type") == "text" and content.get("author") == "user":
                         # Check for user message echo
+                        print("here")
                         if content.get("content") == user_message:
                             user_echo_found = True
 
                 # Exit early if we've found all expected messages
-                if task_creation_found and user_echo_found and agent_response_found:
+                if user_echo_found and agent_response_found:
                     break
 
         # Start streaming task
@@ -160,8 +174,16 @@ async def collect_stream_events():
         event_content = TextContentParam(type="text", author="user", content=user_message)
         await client.agents.send_event(agent_id=agent_id, params={"task_id": task.id, "content": event_content})
 
-        # Wait for streaming to complete
-        await stream_task
+
+        # Wait for the stream to complete (with timeout)
+        try:
+            await asyncio.wait_for(stream_task, timeout=stream_timeout)
+        except asyncio.TimeoutError:
+            pytest.fail(f"Stream timed out after {stream_timeout}s waiting for expected messages")
+
+        # Verify all expected messages were received (fail if stream ended without finding them)
+        assert user_echo_found, "User message echo not found in stream"
+        assert agent_response_found, "Agent response not found in stream"
 
 
 if __name__ == "__main__":

examples/tutorials/10_async/00_base/010_multiturn/tests/test_agent.py

@@ -185,8 +185,11 @@ async def stream_messages():
         event_content = TextContentParam(type="text", author="user", content=user_message)
         await client.agents.send_event(agent_id=agent_id, params={"task_id": task.id, "content": event_content})
 
-        # Wait for streaming to complete
-        await stream_task
+        # Wait for streaming to complete (with timeout)
+        try:
+            await asyncio.wait_for(stream_task, timeout=15)
+        except asyncio.TimeoutError:
+            pytest.fail("Stream timed out after 15s waiting for expected messages")
 
         # Validate we received events
         assert len(all_events) > 0, "No events received in streaming response"