78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
"""SSE 契约测试。"""
|
|
|
|
from collections.abc import AsyncGenerator
|
|
|
|
import pytest
|
|
from httpx import ASGITransport, AsyncClient
|
|
|
|
from app.ai_provider.types.chat import ChatChunk
|
|
from app.application.agent.runtime import AgentTurnResult
|
|
from app.core.config import settings
|
|
from app.main import app
|
|
|
|
|
|
@pytest.fixture
|
|
async def client(
|
|
monkeypatch: pytest.MonkeyPatch, auth_headers: dict[str, str]
|
|
) -> AsyncGenerator[AsyncClient, None]:
|
|
_content = "你好，我是测试模型回复，会根据你的输入生成足够长的内容用于 SSE 分块验证。" * 4
|
|
|
|
async def fake_run_agent_turn(*, on_chunk, **_kwargs: object):
|
|
await on_chunk(ChatChunk(delta_content=_content))
|
|
return AgentTurnResult(content=_content, reasoning="")
|
|
|
|
monkeypatch.setattr("app.application.chat.send_message.run_agent_turn", fake_run_agent_turn)
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(
|
|
transport=transport, base_url="http://test", headers=auth_headers
|
|
) as http_client:
|
|
yield http_client
|
|
|
|
|
|
async def _completed_run(client: AsyncClient, content: str = "你好") -> str:
|
|
created = await client.post("/api/v1/chat/sessions", json={"title": "SSE"})
|
|
session_id = created.json()["session_id"]
|
|
sent = await client.post(
|
|
f"/api/v1/chat/sessions/{session_id}/messages",
|
|
json={"content": content, "modality": "text"},
|
|
)
|
|
return str(sent.json()["run_id"])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completed_run_stream_outputs_events(client: AsyncClient) -> None:
|
|
run_id = await _completed_run(client)
|
|
response = await client.get(f"/api/v1/runs/{run_id}/stream")
|
|
assert response.status_code == 200
|
|
text = response.text
|
|
assert "event: run.started" in text
|
|
assert "event: message.delta" in text
|
|
assert "event: message.completed" in text
|
|
assert "event: run.completed" in text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completed_run_stream_outputs_chunked_message_deltas(client: AsyncClient) -> None:
|
|
run_id = await _completed_run(client, content="你好" * 80)
|
|
response = await client.get(f"/api/v1/runs/{run_id}/stream")
|
|
assert response.status_code == 200
|
|
text = response.text
|
|
assert text.count("event: message.delta") >= 2
|
|
assert "event: message.completed" in text
|
|
assert "event: run.completed" in text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_not_found_and_disabled(
|
|
client: AsyncClient, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
not_found = await client.get("/api/v1/runs/run_missing/stream")
|
|
assert not_found.status_code == 404
|
|
assert not_found.json()["error"]["code"] == "not_found"
|
|
|
|
run_id = await _completed_run(client)
|
|
monkeypatch.setattr(settings, "enable_sse", False)
|
|
disabled = await client.get(f"/api/v1/runs/{run_id}/stream")
|
|
assert disabled.status_code == 400
|
|
assert disabled.json()["error"]["code"] == "invalid_request"
|