temporal

Author: smoreinis

examples/tutorials/10_async/10_temporal/000_hello_acp/tests/test_agent.py

@@ -10,30 +10,94 @@
 """
 
 import pytest
+import pytest_asyncio
 
-from agentex.lib.testing import async_test_agent, assert_valid_agent_response
+from agentex.lib.testing import (
+    async_test_agent,
+    stream_agent_response,
+    assert_valid_agent_response,
+    assert_agent_response_contains,
+)
+from agentex.lib.testing.sessions import AsyncAgentTest
 
 AGENT_NAME = "at000-hello-acp"
 
 
-@pytest.mark.asyncio
-async def test_agent_basic():
-    """Test basic agent functionality."""
-    async with async_test_agent(agent_name=AGENT_NAME) as test:
-        response = await test.send_event("Test message", timeout_seconds=60.0)
+@pytest.fixture
+def agent_name():
+    """Return the agent name for testing."""
+    return AGENT_NAME
+
+
+@pytest_asyncio.fixture
+async def test_agent(agent_name: str):
+    """Fixture to create a test async agent."""
+    async with async_test_agent(agent_name=agent_name) as test:
+        yield test
+
+
+class TestNonStreamingEvents:
+    """Test non-streaming event sending and polling."""
+
+    @pytest.mark.asyncio
+    async def test_send_event_and_poll(self, test_agent: AsyncAgentTest):
+        """Test sending an event and polling for the response."""
+        # Poll for initial task creation message
+        initial_response = await test_agent.poll_for_agent_response(timeout_seconds=15.0)
+        assert_valid_agent_response(initial_response)
+        assert_agent_response_contains(initial_response, "Hello! I've received your task")
+
+        # Send a test message and validate response
+        response = await test_agent.send_event("Hello, this is a test message!", timeout_seconds=30.0)
+        # Validate latest response
         assert_valid_agent_response(response)
+        assert_agent_response_contains(response, "Hello! I've received your message")
+
+
+class TestStreamingEvents:
+    """Test streaming event sending."""
 
+    @pytest.mark.asyncio
+    async def test_send_event_and_stream(self, test_agent: AsyncAgentTest):
+        """Test sending an event and streaming the response."""
+        user_message = "Hello, this is a test message!"
 
-@pytest.mark.asyncio
-async def test_agent_streaming():
-    """Test streaming responses."""
-    async with async_test_agent(agent_name=AGENT_NAME) as test:
-        events = []
-        async for event in test.send_event_and_stream("Stream test", timeout_seconds=60.0):
-            events.append(event)
-            if event.get("type") == "done":
+        # Flags to track what we've received
+        user_echo_found = False
+        agent_response_found = False
+        all_events = []
+
+        # Stream events
+        async for event in stream_agent_response(test_agent.client, test_agent.task_id, timeout=30.0):
+            all_events.append(event)
+            event_type = event.get("type")
+
+            if event_type == "connected":
+                await test_agent.send_event(user_message, timeout_seconds=30.0)
+
+            elif event_type == "full":
+                content = event.get("content", {})
+                if content.get("content") is None:
+                    continue  # Skip empty content
+
+                if content.get("type") == "text" and content.get("author") == "agent":
+                    # Check for agent response to user message
+                    if "Hello! I've received your message" in content.get("content", ""):
+                        agent_response_found = True
+                        assert user_echo_found, "User echo should be found before agent response"
+
+                elif content.get("type") == "text" and content.get("author") == "user":
+                    # Check for user message echo (may or may not be present)
+                    if content.get("content") == user_message:
+                        user_echo_found = True
+
+            # Exit early if we've found expected messages
+            if agent_response_found and user_echo_found:
                 break
-        assert len(events) > 0
+
+        assert agent_response_found, "Did not receive agent response to user message"
+        assert user_echo_found, "User echo message not found"
+        assert len(all_events) > 0, "Should receive events"
 
 
 if __name__ == "__main__":

examples/tutorials/10_async/10_temporal/010_agent_chat/project/workflow.py

@@ -221,7 +221,7 @@ async def on_task_event_send(self, params: SendEventParams) -> None:
                     "to provide accurate and well-reasoned responses."
                 ),
                 parent_span_id=span.id if span else None,
-                model="gpt-4o-mini",
+                model="gpt-5-mini",
                 model_settings=ModelSettings(
                     # Include reasoning items in the response (IDs, summaries)
                     # response_include=["reasoning.encrypted_content"],

examples/tutorials/10_async/10_temporal/010_agent_chat/tests/test_agent.py

@@ -15,183 +15,119 @@
 Run: pytest tests/test_agent.py -v
 """
 
+import asyncio
+
 import pytest
+import pytest_asyncio
 
-from agentex.lib.testing import async_test_agent, assert_valid_agent_response
+from agentex.lib.testing import async_test_agent, stream_agent_response, assert_valid_agent_response
+from agentex.lib.testing.sessions import AsyncAgentTest
+from agentex.types.agent_rpc_result import StreamTaskMessageDone, StreamTaskMessageFull
 
 AGENT_NAME = "at010-agent-chat"
 
 
-@pytest.mark.asyncio
-async def test_agent_basic():
-    """Test basic agent functionality."""
-    async with async_test_agent(agent_name=AGENT_NAME) as test:
-        response = await test.send_event("Test message", timeout_seconds=60.0)
-        assert_valid_agent_response(response)
+@pytest.fixture
+def agent_name():
+    """Return the agent name for testing."""
+    return AGENT_NAME
 
 
-@pytest.mark.asyncio
-async def test_agent_streaming():
-    """Test streaming responses."""
-    async with async_test_agent(agent_name=AGENT_NAME) as test:
-        events = []
-        async for event in test.send_event_and_stream("Stream test", timeout_seconds=60.0):
-            events.append(event)
-            if event.get("type") == "done":
-                break
-        assert len(events) > 0
+@pytest_asyncio.fixture
+async def test_agent(agent_name: str):
+    """Fixture to create a test async agent."""
+    async with async_test_agent(agent_name=agent_name) as test:
+        yield test
+
+class TestNonStreamingEvents:
+    """Test non-streaming event sending and polling with OpenAI Agents SDK."""
 
     @pytest.mark.asyncio
-    async def test_send_event_and_poll_with_calculator(self, client: AsyncAgentex, agent_id: str):
-        """Test sending an event that triggers calculator tool usage and polling for the response."""
-        # Create a task for this conversation
-        task_response = await client.agents.create_task(agent_id, params=ParamsCreateTaskRequest(name=uuid.uuid1().hex))
-        task = task_response.result
-        assert task is not None
+    async def test_send_event_and_poll_simple_query(self, test_agent: AsyncAgentTest):
+        """Test basic agent functionality."""
+        # Wait for state initialization
+        await asyncio.sleep(1)
+
+        # Send a simple message that shouldn't require tool use
+        response = await test_agent.send_event("Hello! Please introduce yourself briefly.", timeout_seconds=30.0)
+        assert_valid_agent_response(response)
 
+    @pytest.mark.asyncio
+    async def test_send_event_and_poll_with_calculator(self, test_agent: AsyncAgentTest):
+        """Test sending an event that triggers calculator tool usage and polling for the response."""
         # Wait for workflow to initialize
         await asyncio.sleep(1)
 
         # Send a message that could trigger the calculator tool (though with reasoning, it may not need it)
         user_message = "What is 15 multiplied by 37?"
-        has_final_agent_response = False
-
-        async for message in send_event_and_poll_yielding(
-            client=client,
-            agent_id=agent_id,
-            task_id=task.id,
-            user_message=user_message,
-            timeout=60,  # Longer timeout for tool use
-            sleep_interval=1.0,
-        ):
-            assert isinstance(message, TaskMessage)
-            if message.content and message.content.type == "text" and message.content.author == "agent":
-                # Check that the answer contains 555 (15 * 37)
-                if "555" in message.content.content:
-                    has_final_agent_response = True
-                    break
-
-        assert has_final_agent_response, "Did not receive final agent text response with correct answer"
+        response = await test_agent.send_event(user_message, timeout_seconds=60.0)
+        assert_valid_agent_response(response)
+        assert "555" in response.content, "Expected answer '555' not found in agent response"
 
     @pytest.mark.asyncio
-    async def test_multi_turn_conversation(self, client: AsyncAgentex, agent_id: str):
+    async def test_multi_turn_conversation_with_state(self, test_agent: AsyncAgentTest):
         """Test multiple turns of conversation with state preservation."""
-        # Create a task for this conversation
-        task_response = await client.agents.create_task(agent_id, params=ParamsCreateTaskRequest(name=uuid.uuid1().hex))
-        task = task_response.result
-        assert task is not None
-
         # Wait for workflow to initialize
         await asyncio.sleep(1)
 
-        # First turn
-        user_message_1 = "My favorite color is blue."
-        async for message in send_event_and_poll_yielding(
-            client=client,
-            agent_id=agent_id,
-            task_id=task.id,
-            user_message=user_message_1,
-            timeout=20,
-            sleep_interval=1.0,
-        ):
-            assert isinstance(message, TaskMessage)
-            if (
-                message.content
-                and message.content.type == "text"
-                and message.content.author == "agent"
-                and message.content.content
-            ):
-                break
-
-        # Wait a bit for state to update
-        await asyncio.sleep(2)
-
-        # Second turn - reference previous context
-        found_response = False
-        user_message_2 = "What did I just tell you my favorite color was?"
-        async for message in send_event_and_poll_yielding(
-            client=client,
-            agent_id=agent_id,
-            task_id=task.id,
-            user_message=user_message_2,
-            timeout=30,
-            sleep_interval=1.0,
-        ):
-            if (
-                message.content
-                and message.content.type == "text"
-                and message.content.author == "agent"
-                and message.content.content
-            ):
-                response_text = message.content.content.lower()
-                assert "blue" in response_text, f"Expected 'blue' in response but got: {response_text}"
-                found_response = True
-                break
-
-        assert found_response, "Did not receive final agent text response with context recall"
+        response = await test_agent.send_event("My favorite color is blue", timeout_seconds=30.0)
+        assert_valid_agent_response(response)
+
+        second_response = await test_agent.send_event(
+            "What did I just tell you my favorite color was?", timeout_seconds=30.0
+        )
+        assert_valid_agent_response(second_response)
+        assert "blue" in second_response.content.lower()
 
 
 class TestStreamingEvents:
     """Test streaming event sending with OpenAI Agents SDK and tool usage."""
 
     @pytest.mark.asyncio
-    async def test_send_event_and_stream_with_reasoning(self, client: AsyncAgentex, agent_id: str):
-        """Test streaming a simple response without tool usage."""
-        # Create a task for this conversation
-        task_response = await client.agents.create_task(agent_id, params=ParamsCreateTaskRequest(name=uuid.uuid1().hex))
-        task = task_response.result
-        assert task is not None
-
+    async def test_send_event_and_stream_with_reasoning(self, test_agent: AsyncAgentTest):
+        """Test streaming event responses."""
         # Wait for workflow to initialize
         await asyncio.sleep(1)
 
+        # Send message and stream response
         user_message = "Tell me a very short joke about programming."
 
         # Check for user message and agent response
         user_message_found = False
         agent_response_found = False
 
-        async def stream_messages() -> None:  # noqa: ANN101
-            nonlocal user_message_found, agent_response_found
-            async for event in stream_agent_response(
-                client=client,
-                task_id=task.id,
-                timeout=60,
-            ):
-                msg_type = event.get("type")
-                if msg_type == "full":
-                    task_message_update = StreamTaskMessageFull.model_validate(event)
-                    if task_message_update.parent_task_message and task_message_update.parent_task_message.id:
-                        finished_message = await client.messages.retrieve(task_message_update.parent_task_message.id)
-                        if (
-                            finished_message.content
-                            and finished_message.content.type == "text"
-                            and finished_message.content.author == "user"
-                        ):
-                            user_message_found = True
-                        elif (
-                            finished_message.content
-                            and finished_message.content.type == "text"
-                            and finished_message.content.author == "agent"
-                        ):
-                            agent_response_found = True
-                        elif finished_message.content and finished_message.content.type == "reasoning":
-                            tool_response_found = True
-                elif msg_type == "done":
-                    task_message_update = StreamTaskMessageDone.model_validate(event)
-                    if task_message_update.parent_task_message and task_message_update.parent_task_message.id:
-                        finished_message = await client.messages.retrieve(task_message_update.parent_task_message.id)
-                        if finished_message.content and finished_message.content.type == "reasoning":
-                            agent_response_found = True
-                    continue
-
-        stream_task = asyncio.create_task(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
+        # Stream events
+        async for event in stream_agent_response(test_agent.client, test_agent.task_id, timeout=60.0):
+            event_type = event.get("type")
+
+            if event_type == "connected":
+                await test_agent.send_event(user_message, timeout_seconds=30.0)
+
+            elif event_type == "full":
+                print('full event', event)
+                task_message_update = StreamTaskMessageFull.model_validate(event)
+                if task_message_update.parent_task_message and task_message_update.parent_task_message.id:
+                    finished_message = await test_agent.client.messages.retrieve(task_message_update.parent_task_message.id)
+                    if (
+                        finished_message.content
+                        and finished_message.content.type == "text"
+                        and finished_message.content.author == "user"
+                    ):
+                        user_message_found = True
+                    elif (
+                        finished_message.content
+                        and finished_message.content.type == "text"
+                        and finished_message.content.author == "agent"
+                    ):
+                        agent_response_found = True
+            elif event_type == "done":
+                print('done event', event)
+                task_message_update = StreamTaskMessageDone.model_validate(event)
+                if task_message_update.parent_task_message and task_message_update.parent_task_message.id:
+                    finished_message = await test_agent.client.messages.retrieve(task_message_update.parent_task_message.id)
+                    if finished_message.content and finished_message.content.type == "text" and finished_message.content.author == "agent":
+                        agent_response_found = True
+                continue
 
         assert user_message_found, "User message not found in stream"
         assert agent_response_found, "Agent response not found in stream"