mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-17 18:25:42 +02:00
70 lines
3.1 KiB
Python
70 lines
3.1 KiB
Python
"""Direct coverage for the extracted AssistantMessage handler. The harness drives the common
|
|
paths; these pin the branches it doesn't: the router-auth-expiry friendly card, tool-call
|
|
commit, and per-turn output-token accumulation."""
|
|
|
|
import pytest
|
|
from unittest.mock import patch, AsyncMock
|
|
|
|
from claude_agent_sdk import AssistantMessage
|
|
from claude_agent_sdk.types import TextBlock, ToolUseBlock
|
|
|
|
from backend.apps.agents.core.models import AgentSession
|
|
from backend.apps.agents.manager.streaming.state import TurnState, ThinkingState
|
|
from backend.apps.agents.manager.streaming import handle_assistant_message as assistant_message
|
|
|
|
|
|
def p_asst(blocks, usage=None):
|
|
return AssistantMessage(content=blocks, model="sonnet", message_id="m1", stop_reason="end_turn",
|
|
session_id="s", usage=usage or {"input_tokens": 1, "output_tokens": 1})
|
|
|
|
|
|
def p_fixt():
|
|
return AgentSession(name="t", model="sonnet", dashboard_id="d"), TurnState(), ThinkingState()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_plain_text_commits_assistant_message():
|
|
session, turn, thinking = p_fixt()
|
|
with patch.object(assistant_message.ws_manager, "send_to_session", new=AsyncMock()):
|
|
await assistant_message.handle_assistant_message(
|
|
p_asst([TextBlock(text="Hello there")]), session, session.id, turn, thinking, {}, {})
|
|
assert any(m.role == "assistant" and "Hello there" in str(m.content) for m in session.messages)
|
|
assert turn.number == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_auth_error_surfaces_card_not_assistant_text():
|
|
session, turn, thinking = p_fixt()
|
|
txt = "[codex/gpt-5] Failed to authenticate: 401 provided authentication token is expired"
|
|
events = []
|
|
|
|
async def fake_send(sid, event, data):
|
|
events.append(event)
|
|
|
|
with patch.object(assistant_message.ws_manager, "send_to_session", new=fake_send):
|
|
await assistant_message.handle_assistant_message(
|
|
p_asst([TextBlock(text=txt)]), session, session.id, turn, thinking, {}, {})
|
|
assert any(m.role == "system" for m in session.messages) # friendly card
|
|
assert not any(m.role == "assistant" for m in session.messages) # NOT committed as the reply
|
|
assert "agent:auth_error" in events
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_use_block_commits_tool_call():
|
|
session, turn, thinking = p_fixt()
|
|
with patch.object(assistant_message.ws_manager, "send_to_session", new=AsyncMock()):
|
|
await assistant_message.handle_assistant_message(
|
|
p_asst([ToolUseBlock(id="tu1", name="Read", input={"file_path": "/x"})]),
|
|
session, session.id, turn, thinking, {}, {})
|
|
assert any(m.role == "tool_call" for m in session.messages)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_tokens_accumulate_onto_turn():
|
|
session, turn, thinking = p_fixt()
|
|
with patch.object(assistant_message.ws_manager, "send_to_session", new=AsyncMock()):
|
|
await assistant_message.handle_assistant_message(
|
|
p_asst([TextBlock(text="hi")], usage={"input_tokens": 10, "output_tokens": 42}),
|
|
session, session.id, turn, thinking, {}, {})
|
|
assert turn.output_tokens == 42
|