diff --git a/examples/tutorials/10_async/00_base/010_multiturn/tests/test_agent.py b/examples/tutorials/10_async/00_base/010_multiturn/tests/test_agent.py index 43d283b8c..60bad474e 100644 --- a/examples/tutorials/10_async/00_base/010_multiturn/tests/test_agent.py +++ b/examples/tutorials/10_async/00_base/010_multiturn/tests/test_agent.py @@ -24,6 +24,7 @@ import pytest_asyncio from test_utils.async_utils import ( stream_agent_response, + wait_for_state_messages, send_event_and_poll_yielding, ) @@ -124,11 +125,7 @@ async def test_send_event_and_poll(self, client: AsyncAgentex, agent_id: str): assert user_message_found, "User message not found" assert agent_response_found, "Agent response not found" - await asyncio.sleep(1) # wait for state to be updated - states = await client.states.list(agent_id=agent_id, task_id=task.id) - assert len(states) == 1 - state = states[0].state - messages = state.get("messages", []) + messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3) assert isinstance(messages, list) assert len(messages) == 3 @@ -207,11 +204,7 @@ async def stream_messages() -> None: assert agent_response_found, "Agent response not found in stream" # Verify the state has been updated - await asyncio.sleep(1) # wait for state to be updated - states = await client.states.list(agent_id=agent_id, task_id=task.id) - assert len(states) == 1 - state = states[0].state - messages = state.get("messages", []) + messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3) assert isinstance(messages, list) assert len(messages) == 3 diff --git a/examples/tutorials/10_async/00_base/020_streaming/tests/test_agent.py b/examples/tutorials/10_async/00_base/020_streaming/tests/test_agent.py index c55525191..d64acbca8 100644 --- a/examples/tutorials/10_async/00_base/020_streaming/tests/test_agent.py +++ b/examples/tutorials/10_async/00_base/020_streaming/tests/test_agent.py @@ -24,6 +24,7 @@ import pytest_asyncio from test_utils.async_utils import ( stream_agent_response, + wait_for_state_messages, send_event_and_poll_yielding, ) @@ -122,11 +123,7 @@ async def test_send_event_and_poll(self, client: AsyncAgentex, agent_id: str): assert agent_response_found, "Agent response not found" # assert the state has been updated - await asyncio.sleep(1) # wait for state to be updated - states = await client.states.list(agent_id=agent_id, task_id=task.id) - assert len(states) == 1 - state = states[0].state - messages = state.get("messages", []) + messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3) assert isinstance(messages, list) assert len(messages) == 3 @@ -205,11 +202,7 @@ async def stream_messages() -> None: assert delta_messages_found, "Delta messages not found in stream (streaming response expected)" # Verify the state has been updated - await asyncio.sleep(1) # wait for state to be updated - states = await client.states.list(agent_id=agent_id, task_id=task.id) - assert len(states) == 1 - state: dict[str, object] = states[0].state - messages = state.get("messages", []) + messages = await wait_for_state_messages(client, agent_id, task.id, expected_count=3) assert isinstance(messages, list) assert len(messages) == 3 diff --git a/examples/tutorials/test_utils/async_utils.py b/examples/tutorials/test_utils/async_utils.py index 2187e98d8..cb5e10d3f 100644 --- a/examples/tutorials/test_utils/async_utils.py +++ b/examples/tutorials/test_utils/async_utils.py @@ -156,6 +156,42 @@ async def poll_messages( await asyncio.sleep(sleep_interval) +async def wait_for_state_messages( + client: AsyncAgentex, + agent_id: str, + task_id: str, + expected_count: int, + timeout: float = 30, + sleep_interval: float = 0.5, +) -> list: + """ + Poll the task state until its "messages" list reaches expected_count. + + Agents typically send their reply before they persist the turn to task state, + so reading the state right after the reply arrives can see the previous turn. + + Args: + client: AgentEx client instance + agent_id: The agent ID + task_id: The task ID + expected_count: Number of state messages to wait for + timeout: Maximum seconds to poll (default: 30) + sleep_interval: Seconds to sleep between polls (default: 0.5) + + Returns: + The state's "messages" list from the last poll. Callers assert on it, so a + state that never reaches expected_count still fails with the real length. + """ + deadline = time.monotonic() + timeout + while True: + states = await client.states.list(agent_id=agent_id, task_id=task_id) + assert len(states) == 1, f"Expected exactly one state, got {len(states)}" + messages = states[0].state.get("messages", []) + if len(messages) >= expected_count or time.monotonic() >= deadline: + return messages + await asyncio.sleep(sleep_interval) + + async def send_event_and_stream( client: AsyncAgentex, agent_id: str, diff --git a/tests/lib/test_tutorial_async_utils.py b/tests/lib/test_tutorial_async_utils.py new file mode 100644 index 000000000..867f01868 --- /dev/null +++ b/tests/lib/test_tutorial_async_utils.py @@ -0,0 +1,79 @@ +"""Tests for wait_for_state_messages in the tutorials' test_utils. + +The tutorials' `test_utils` package shares its name with `tests/test_utils`, so the +module is loaded from its file path instead of imported by name. +""" + +from __future__ import annotations + +import time +import importlib.util +from types import ModuleType, SimpleNamespace +from pathlib import Path + +import pytest + +_ASYNC_UTILS = Path(__file__).resolve().parents[2] / "examples" / "tutorials" / "test_utils" / "async_utils.py" + + +def _load_async_utils() -> ModuleType: + spec = importlib.util.spec_from_file_location("tutorial_async_utils", _ASYNC_UTILS) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +wait_for_state_messages = _load_async_utils().wait_for_state_messages + + +class _FakeStates: + """Returns one state whose message count grows from 1 to 3 once `write_after` seconds pass.""" + + def __init__(self, write_after: float | None): + self._start = time.monotonic() + self._write_after = write_after + self.calls = 0 + + async def list(self, agent_id: str, task_id: str) -> list[SimpleNamespace]: + del agent_id, task_id + self.calls += 1 + written = self._write_after is not None and time.monotonic() - self._start >= self._write_after + return [SimpleNamespace(state={"messages": [{"role": "system"}] * (3 if written else 1)})] + + +def _client(write_after: float | None) -> SimpleNamespace: + return SimpleNamespace(states=_FakeStates(write_after)) + + +@pytest.mark.asyncio +async def test_waits_for_a_state_write_that_lands_after_the_reply() -> None: + client = _client(write_after=0.3) + + messages = await wait_for_state_messages(client, "agent", "task", expected_count=3, timeout=5, sleep_interval=0.05) + + assert len(messages) == 3 + assert client.states.calls > 1 + + +@pytest.mark.asyncio +async def test_returns_the_real_state_at_the_deadline_when_it_never_updates() -> None: + client = _client(write_after=None) + start = time.monotonic() + + messages = await wait_for_state_messages(client, "agent", "task", expected_count=3, timeout=0.3, sleep_interval=0.05) + + # It kept polling until the deadline instead of returning early. No upper bound: + # a busy worker may wake late, and returning at all already proves it stopped. + assert len(messages) == 1 + assert time.monotonic() - start >= 0.3 + + +@pytest.mark.asyncio +async def test_returns_on_the_first_poll_when_the_state_is_already_written() -> None: + client = _client(write_after=0) + + messages = await wait_for_state_messages(client, "agent", "task", expected_count=3) + + assert len(messages) == 3 + assert client.states.calls == 1