289 lines
9.6 KiB
Python
289 lines
9.6 KiB
Python
"""Integration tests for the SSE report-events endpoint (GET /reports/{id}/events).
|
||
|
||
Run against the Docker Compose infrastructure (`docker compose up -d`).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
import pytest
|
||
from sqlalchemy import text
|
||
|
||
from tests.integration.conftest import user_token
|
||
|
||
pytestmark = pytest.mark.integration
|
||
|
||
_FINDINGS = [
|
||
{
|
||
"checklist_id": "penalties",
|
||
"severity": "high",
|
||
"quote": "Неустойка 10% от суммы договора",
|
||
"section_ref": "5.2",
|
||
"risk": "Односторонняя неустойка",
|
||
"recommendation": "Согласовать взаимную ответственность",
|
||
},
|
||
{
|
||
"checklist_id": "jurisdiction",
|
||
"severity": "medium",
|
||
"quote": "Споры рассматриваются в суде г. Минск",
|
||
"section_ref": "9.1",
|
||
"risk": "Неудобная подсудность",
|
||
"recommendation": "",
|
||
},
|
||
]
|
||
|
||
|
||
@pytest.fixture
|
||
def pdf_bytes(tmp_path: Path) -> bytes:
|
||
import pymupdf
|
||
|
||
doc = pymupdf.open()
|
||
page = doc.new_page()
|
||
page.insert_text((72, 72), "Договор. Стороны обязуются.")
|
||
path = tmp_path / "contract.pdf"
|
||
doc.save(str(path))
|
||
doc.close()
|
||
return path.read_bytes()
|
||
|
||
|
||
async def _create_user_with_credits(db_session, telegram_id: int) -> None: # noqa: ANN001
|
||
await db_session.execute(
|
||
text(
|
||
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, 10) "
|
||
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 10"
|
||
),
|
||
{"t": telegram_id},
|
||
)
|
||
await db_session.commit()
|
||
|
||
|
||
async def _upload(client: httpx.AsyncClient, token: str, pdf_bytes: bytes) -> str:
|
||
"""Upload a document under user JWT and return its id."""
|
||
resp = await client.post(
|
||
"/api/v1/documents",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
files={"file": ("contract.pdf", pdf_bytes, "application/pdf")},
|
||
)
|
||
assert resp.status_code in (200, 202), resp.text
|
||
return resp.json()["document_id"]
|
||
|
||
|
||
def _parse_events(raw: str) -> list[tuple[str, dict[str, Any]]]:
|
||
"""Parse an SSE body into (event, data) tuples, ignoring comment lines."""
|
||
events: list[tuple[str, dict[str, Any]]] = []
|
||
for chunk in raw.split("\n\n"):
|
||
chunk = chunk.strip()
|
||
if not chunk or chunk.startswith(":"):
|
||
continue
|
||
event = "message"
|
||
data = ""
|
||
for line in chunk.split("\n"):
|
||
if line.startswith("event:"):
|
||
event = line[len("event:") :].strip()
|
||
elif line.startswith("data:"):
|
||
data = line[len("data:") :].strip()
|
||
if data:
|
||
events.append((event, json.loads(data)))
|
||
return events
|
||
|
||
|
||
async def _first_events(
|
||
client: httpx.AsyncClient,
|
||
token: str,
|
||
document_id: str,
|
||
count: int,
|
||
) -> list[tuple[str, dict[str, Any]]]:
|
||
"""Open the SSE stream and collect the first `count` parsed events."""
|
||
collected: list[tuple[str, dict[str, Any]]] = []
|
||
async with client.stream(
|
||
"GET",
|
||
f"/api/v1/reports/{document_id}/events",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
) as response:
|
||
assert response.status_code == 200
|
||
assert response.headers["content-type"].startswith("text/event-stream")
|
||
buffer = ""
|
||
async for chunk in response.aiter_text():
|
||
buffer += chunk
|
||
collected = _parse_events(buffer)
|
||
if len(collected) >= count:
|
||
break
|
||
return collected[:count]
|
||
|
||
|
||
async def test_sse_requires_auth(client: httpx.AsyncClient) -> None:
|
||
response = await client.get(f"/api/v1/reports/{uuid.uuid4()}/events")
|
||
assert response.status_code == 401
|
||
|
||
|
||
async def test_sse_unknown_document_returns_404(
|
||
client: httpx.AsyncClient,
|
||
infra: dict[str, str],
|
||
) -> None:
|
||
telegram_id = 999100101
|
||
token = await user_token(client, infra, telegram_id)
|
||
response = await client.get(
|
||
f"/api/v1/reports/{uuid.uuid4()}/events",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
)
|
||
assert response.status_code == 404
|
||
|
||
|
||
async def test_sse_streams_progress_then_completed(
|
||
client: httpx.AsyncClient,
|
||
db_session, # noqa: ANN001
|
||
infra: dict[str, str],
|
||
pdf_bytes: bytes,
|
||
) -> None:
|
||
telegram_id = 999100102
|
||
await _create_user_with_credits(db_session, telegram_id)
|
||
token = await user_token(client, infra, telegram_id)
|
||
document_id = await _upload(client, token, pdf_bytes)
|
||
|
||
# Simulate the pipeline finishing while the client is subscribed: the
|
||
# stream should observe the row change on its next poll tick.
|
||
async def _finish() -> None:
|
||
await asyncio.sleep(1.5)
|
||
await db_session.execute(
|
||
text(
|
||
"INSERT INTO reports (document_id, content_json, markdown, model_used) "
|
||
"VALUES (:d, :cj, :md, 'test-model')"
|
||
),
|
||
{
|
||
"d": document_id,
|
||
"cj": json.dumps({"findings": _FINDINGS}),
|
||
"md": "# Отчёт",
|
||
},
|
||
)
|
||
await db_session.execute(
|
||
text("UPDATE documents SET status = 'done' WHERE id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
await db_session.commit()
|
||
|
||
task = asyncio.create_task(_finish())
|
||
try:
|
||
events = await _first_events(client, token, document_id, 2)
|
||
finally:
|
||
await task
|
||
|
||
(first_name, first) = events[0]
|
||
(last_name, last) = events[-1]
|
||
assert first_name == last_name == "message"
|
||
|
||
# Progress frame: analysisResultSchema with empty result fields.
|
||
assert first["id"] == document_id
|
||
assert first["fileName"] == "contract.pdf"
|
||
assert first["fileType"] == "application/pdf"
|
||
assert first["fileSize"] == len(pdf_bytes)
|
||
assert first["status"] in ("pending", "processing")
|
||
assert first["issues"] == []
|
||
assert first["riskScore"] == 0
|
||
assert first["completedAt"] is None
|
||
assert first["createdAt"]
|
||
|
||
# Terminal frame: full mapped result.
|
||
assert last["status"] == "completed"
|
||
assert last["completedAt"] is not None
|
||
assert last["riskScore"] == 6 # critical(4) + warning(2)
|
||
assert last["summary"] == "Критичных: 1, предупреждений: 1, замечаний: 0"
|
||
|
||
issues = last["issues"]
|
||
assert len(issues) == 2
|
||
assert issues[0]["id"] == "penalties-0"
|
||
assert issues[0]["severity"] == "critical"
|
||
assert issues[0]["category"] == "penalties"
|
||
assert issues[0]["title"] == "Неустойки / штрафы"
|
||
assert "Односторонняя неустойка" in issues[0]["description"]
|
||
assert "Рекомендация: Согласовать" in issues[0]["description"]
|
||
assert issues[0]["fragment"] == "Неустойка 10% от суммы договора (п. 5.2)"
|
||
assert issues[0]["lineNumber"] is None
|
||
assert issues[1]["severity"] == "warning"
|
||
assert "Рекомендация" not in issues[1]["description"]
|
||
|
||
|
||
async def test_sse_failed_document_emits_failed_and_closes(
|
||
client: httpx.AsyncClient,
|
||
db_session, # noqa: ANN001
|
||
infra: dict[str, str],
|
||
pdf_bytes: bytes,
|
||
) -> None:
|
||
telegram_id = 999100103
|
||
await _create_user_with_credits(db_session, telegram_id)
|
||
token = await user_token(client, infra, telegram_id)
|
||
document_id = await _upload(client, token, pdf_bytes)
|
||
|
||
await db_session.execute(
|
||
text("UPDATE documents SET status = 'failed', stage = 'analyze' WHERE id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
await db_session.commit()
|
||
|
||
events = await _first_events(client, token, document_id, 1)
|
||
name, payload = events[0]
|
||
assert name == "message"
|
||
assert payload["id"] == document_id
|
||
assert payload["status"] == "failed"
|
||
assert payload["issues"] == []
|
||
assert payload["completedAt"] is None
|
||
|
||
|
||
async def test_sse_manual_review_maps_to_completed(
|
||
client: httpx.AsyncClient,
|
||
db_session, # noqa: ANN001
|
||
infra: dict[str, str],
|
||
pdf_bytes: bytes,
|
||
) -> None:
|
||
telegram_id = 999100105
|
||
await _create_user_with_credits(db_session, telegram_id)
|
||
token = await user_token(client, infra, telegram_id)
|
||
document_id = await _upload(client, token, pdf_bytes)
|
||
|
||
await db_session.execute(
|
||
text("UPDATE documents SET status = 'manual_review' WHERE id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
await db_session.commit()
|
||
|
||
events = await _first_events(client, token, document_id, 1)
|
||
_name, payload = events[0]
|
||
assert payload["status"] == "completed"
|
||
assert payload["issues"] == []
|
||
assert payload["riskScore"] == 0
|
||
assert payload["summary"] == ""
|
||
|
||
|
||
async def test_sse_emits_keepalive_comments_while_unchanged(
|
||
client: httpx.AsyncClient,
|
||
db_session, # noqa: ANN001
|
||
infra: dict[str, str],
|
||
pdf_bytes: bytes,
|
||
) -> None:
|
||
"""While the status does not change, only comment lines arrive."""
|
||
telegram_id = 999100104
|
||
await _create_user_with_credits(db_session, telegram_id)
|
||
token = await user_token(client, infra, telegram_id)
|
||
document_id = await _upload(client, token, pdf_bytes)
|
||
|
||
raw = ""
|
||
async with client.stream(
|
||
"GET",
|
||
f"/api/v1/reports/{document_id}/events",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
) as response:
|
||
assert response.status_code == 200
|
||
async for chunk in response.aiter_text():
|
||
raw += chunk
|
||
if raw.count(": keep-alive") >= 2:
|
||
break
|
||
|
||
assert "data: " in raw
|
||
assert raw.count(": keep-alive") >= 2
|
||
# Only one data frame (the initial snapshot) while nothing changes.
|
||
assert len(_parse_events(raw)) == 1
|