forked from Rockachopa/Timmy-time-dashboard
Compare commits
1 Commits
kimi/issue
...
kimi/issue
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9ec82ab6ad |
281
tests/timmy/test_agentic_loop_integration.py
Normal file
281
tests/timmy/test_agentic_loop_integration.py
Normal file
@@ -0,0 +1,281 @@
|
||||
"""Integration tests for agentic loop WebSocket broadcasts.
|
||||
|
||||
Verifies that agentic loop events (plan_ready, step_complete, task_complete)
|
||||
flow through the real WebSocketManager to connected clients, and that the
|
||||
loop survives a WS disconnect mid-execution.
|
||||
|
||||
Fixes #445
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from infrastructure.ws_manager.handler import WebSocketManager
|
||||
from timmy.agentic_loop import run_agentic_loop
|
||||
|
||||
|
||||
def _mock_run(content: str):
|
||||
"""Create a mock return value for agent.run()."""
|
||||
m = MagicMock()
|
||||
m.content = content
|
||||
return m
|
||||
|
||||
|
||||
def _fake_ws() -> AsyncMock:
|
||||
"""Return an AsyncMock that behaves like a WebSocket connection."""
|
||||
ws = AsyncMock()
|
||||
ws.send_text = AsyncMock()
|
||||
return ws
|
||||
|
||||
|
||||
def _collect_events(ws_mock: AsyncMock) -> list[dict]:
|
||||
"""Extract parsed JSON events from a mock WebSocket's send_text calls."""
|
||||
events = []
|
||||
for call in ws_mock.send_text.call_args_list:
|
||||
raw = call.args[0] if call.args else call.kwargs.get("data", "")
|
||||
events.append(json.loads(raw))
|
||||
return events
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration: events reach a real WebSocketManager
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plan_ready_event_received_by_ws_client():
|
||||
"""A connected WS client receives an agentic.plan_ready event."""
|
||||
mgr = WebSocketManager()
|
||||
ws = _fake_ws()
|
||||
mgr._connections = [ws]
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. Gather data\n2. Summarise"),
|
||||
_mock_run("Data gathered"),
|
||||
_mock_run("Summary written"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
await run_agentic_loop("Gather and summarise", max_steps=2)
|
||||
|
||||
events = _collect_events(ws)
|
||||
plan_events = [e for e in events if e["event"] == "agentic.plan_ready"]
|
||||
assert len(plan_events) == 1
|
||||
assert plan_events[0]["data"]["steps"] == ["Gather data", "Summarise"]
|
||||
assert plan_events[0]["data"]["total"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_step_complete_events_received_by_ws_client():
|
||||
"""A connected WS client receives agentic.step_complete for each step."""
|
||||
mgr = WebSocketManager()
|
||||
ws = _fake_ws()
|
||||
mgr._connections = [ws]
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. Step A\n2. Step B"),
|
||||
_mock_run("A done"),
|
||||
_mock_run("B done"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
await run_agentic_loop("Do A and B", max_steps=2)
|
||||
|
||||
events = _collect_events(ws)
|
||||
step_events = [e for e in events if e["event"] == "agentic.step_complete"]
|
||||
assert len(step_events) == 2
|
||||
assert step_events[0]["data"]["step"] == 1
|
||||
assert step_events[1]["data"]["step"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_task_complete_event_received_by_ws_client():
|
||||
"""A connected WS client receives agentic.task_complete at the end."""
|
||||
mgr = WebSocketManager()
|
||||
ws = _fake_ws()
|
||||
mgr._connections = [ws]
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. Only step"),
|
||||
_mock_run("Done"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
result = await run_agentic_loop("One step task", max_steps=1)
|
||||
|
||||
events = _collect_events(ws)
|
||||
task_events = [e for e in events if e["event"] == "agentic.task_complete"]
|
||||
assert len(task_events) == 1
|
||||
assert task_events[0]["data"]["status"] == "completed"
|
||||
assert task_events[0]["data"]["task_id"] == result.task_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_event_types_received_in_order():
|
||||
"""Client receives plan_ready → step_complete(s) → task_complete in order."""
|
||||
mgr = WebSocketManager()
|
||||
ws = _fake_ws()
|
||||
mgr._connections = [ws]
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. First\n2. Second"),
|
||||
_mock_run("First done"),
|
||||
_mock_run("Second done"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
await run_agentic_loop("Two steps", max_steps=2)
|
||||
|
||||
events = _collect_events(ws)
|
||||
event_names = [e["event"] for e in events]
|
||||
assert event_names == [
|
||||
"agentic.plan_ready",
|
||||
"agentic.step_complete",
|
||||
"agentic.step_complete",
|
||||
"agentic.task_complete",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_adapted_step_broadcasts_step_adapted():
|
||||
"""When a step fails and adapts, client receives agentic.step_adapted."""
|
||||
mgr = WebSocketManager()
|
||||
ws = _fake_ws()
|
||||
mgr._connections = [ws]
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. Risky step"),
|
||||
Exception("network error"),
|
||||
_mock_run("Adapted approach worked"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
await run_agentic_loop("Risky task", max_steps=1)
|
||||
|
||||
events = _collect_events(ws)
|
||||
event_names = [e["event"] for e in events]
|
||||
assert "agentic.step_adapted" in event_names
|
||||
adapted = [e for e in events if e["event"] == "agentic.step_adapted"][0]
|
||||
assert adapted["data"]["error"] == "network error"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Resilience: WS disconnect mid-loop
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_disconnect_mid_loop_does_not_crash():
|
||||
"""If a WS client disconnects during the loop, the loop completes."""
|
||||
mgr = WebSocketManager()
|
||||
ws = _fake_ws()
|
||||
# First send succeeds (plan_ready), then WS dies
|
||||
ws.send_text = AsyncMock(side_effect=[None, ConnectionError("gone")])
|
||||
mgr._connections = [ws]
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. Step A\n2. Step B"),
|
||||
_mock_run("A done"),
|
||||
_mock_run("B done"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
result = await run_agentic_loop("Survive disconnect", max_steps=2)
|
||||
|
||||
# Loop completed despite WS failure
|
||||
assert result.status == "completed"
|
||||
assert len(result.steps) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_ws_connections_does_not_crash():
|
||||
"""Loop completes normally when no WS clients are connected."""
|
||||
mgr = WebSocketManager()
|
||||
# No connections at all
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. Solo step"),
|
||||
_mock_run("Done alone"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
result = await run_agentic_loop("No audience", max_steps=1)
|
||||
|
||||
assert result.status == "completed"
|
||||
assert len(result.steps) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_ws_clients_all_receive_events():
|
||||
"""All connected WS clients receive the same broadcast events."""
|
||||
mgr = WebSocketManager()
|
||||
ws1 = _fake_ws()
|
||||
ws2 = _fake_ws()
|
||||
mgr._connections = [ws1, ws2]
|
||||
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.run = MagicMock(
|
||||
side_effect=[
|
||||
_mock_run("1. Shared step"),
|
||||
_mock_run("Shared result"),
|
||||
]
|
||||
)
|
||||
|
||||
with (
|
||||
patch("timmy.agentic_loop._get_loop_agent", return_value=mock_agent),
|
||||
patch("infrastructure.ws_manager.handler.ws_manager", mgr),
|
||||
):
|
||||
await run_agentic_loop("Broadcast to all", max_steps=1)
|
||||
|
||||
events1 = _collect_events(ws1)
|
||||
events2 = _collect_events(ws2)
|
||||
|
||||
names1 = [e["event"] for e in events1]
|
||||
names2 = [e["event"] for e in events2]
|
||||
assert names1 == names2
|
||||
assert "agentic.plan_ready" in names1
|
||||
assert "agentic.task_complete" in names1
|
||||
Reference in New Issue
Block a user