101 lines
3.8 KiB
Python
101 lines
3.8 KiB
Python
"""幂等键测试。"""
|
|
|
|
from collections.abc import AsyncGenerator
|
|
|
|
import pytest
|
|
from httpx import ASGITransport, AsyncClient
|
|
from sqlalchemy import select
|
|
|
|
from app.ai_provider.types.chat import ChatChunk
|
|
from app.application.agent.runtime import AgentTurnResult
|
|
from app.application.chat.send_message import drain_pending_streams
|
|
from app.db.session import AsyncSessionLocal
|
|
from app.main import app
|
|
from app.models.chat import ChatMessage
|
|
from app.models.idempotency import IdempotencyKey
|
|
from app.models.memory import Memory
|
|
from app.models.runtime import Run
|
|
|
|
|
|
@pytest.fixture
|
|
async def client(
|
|
monkeypatch: pytest.MonkeyPatch, auth_headers: dict[str, str]
|
|
) -> AsyncGenerator[AsyncClient, None]:
|
|
async def fake_run_agent_turn(*, on_chunk, **_kwargs: object):
|
|
await on_chunk(ChatChunk(delta_content="你好，我是测试模型回复。"))
|
|
return AgentTurnResult(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 _create_session(client: AsyncClient) -> str:
|
|
response = await client.post("/api/v1/chat/sessions", json={"title": "幂等测试"})
|
|
assert response.status_code == 200
|
|
return str(response.json()["session_id"])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_send_message_idempotency_replays_same_response(client: AsyncClient) -> None:
|
|
session_id = await _create_session(client)
|
|
headers = {"Idempotency-Key": "send-same-1"}
|
|
payload = {"content": "你好", "modality": "text"}
|
|
first = await client.post(
|
|
f"/api/v1/chat/sessions/{session_id}/messages", json=payload, headers=headers
|
|
)
|
|
second = await client.post(
|
|
f"/api/v1/chat/sessions/{session_id}/messages", json=payload, headers=headers
|
|
)
|
|
assert first.status_code == 200
|
|
assert second.status_code == 200
|
|
assert second.json() == first.json()
|
|
await drain_pending_streams()
|
|
|
|
async with AsyncSessionLocal() as db:
|
|
runs = (await db.execute(select(Run))).scalars().all()
|
|
messages = (await db.execute(select(ChatMessage))).scalars().all()
|
|
keys = (await db.execute(select(IdempotencyKey))).scalars().all()
|
|
assert len(runs) == 1
|
|
assert len(messages) == 2
|
|
assert len(keys) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_idempotency_key_conflict(client: AsyncClient) -> None:
|
|
session_id = await _create_session(client)
|
|
headers = {"Idempotency-Key": "send-conflict-1"}
|
|
first = await client.post(
|
|
f"/api/v1/chat/sessions/{session_id}/messages",
|
|
json={"content": "你好", "modality": "text"},
|
|
headers=headers,
|
|
)
|
|
conflict = await client.post(
|
|
f"/api/v1/chat/sessions/{session_id}/messages",
|
|
json={"content": "不同内容", "modality": "text"},
|
|
headers=headers,
|
|
)
|
|
assert first.status_code == 200
|
|
assert conflict.status_code == 409
|
|
assert conflict.json()["error"]["code"] == "conflict"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_create_idempotency(client: AsyncClient) -> None:
|
|
headers = {"Idempotency-Key": "memory-same-1"}
|
|
payload = {"type": "user", "title": "偏好", "content": "用户喜欢简洁回答。"}
|
|
first = await client.post("/api/v1/memory", json=payload, headers=headers)
|
|
second = await client.post("/api/v1/memory", json=payload, headers=headers)
|
|
assert first.status_code == 200
|
|
assert second.status_code == 200
|
|
assert second.json()["id"] == first.json()["id"]
|
|
|
|
async with AsyncSessionLocal() as db:
|
|
memories = (await db.execute(select(Memory))).scalars().all()
|
|
assert len(memories) == 1
|