368 lines
13 KiB
Python
368 lines
13 KiB
Python
"""Integration tests for manual-review queue actions (issue 018).
|
||
|
||
These are the "worker handler seam" tests: actions are driven in-process against
|
||
real Postgres / RabbitMQ / MinIO, and the analyze handler is exercised with a stub
|
||
LLM provider to prove that ``send_to_analysis`` produces a real Report.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from decimal import Decimal
|
||
from typing import Any
|
||
|
||
import pytest
|
||
from sqlalchemy import text
|
||
|
||
from contract_check.core.analysis.report_schema import Finding
|
||
from contract_check.core.db.repositories import (
|
||
CreditsRepository,
|
||
DocumentRepository,
|
||
ReportRepository,
|
||
)
|
||
from contract_check.core.db.session import create_session_factory
|
||
from contract_check.core.llm.port import AnalysisResult
|
||
from contract_check.core.mq.messages import AnalyzeRequested
|
||
from contract_check.core.mq.topology import RK_ANALYZE
|
||
from contract_check.core.review.actions import (
|
||
ReviewActionError,
|
||
complete,
|
||
reject,
|
||
send_to_analysis,
|
||
)
|
||
from contract_check.core.s3 import extracted_key
|
||
from contract_check.core.s3.minio_storage import MinioStorage
|
||
from contract_check.worker_analyze.handler import AnalyzeHandler
|
||
|
||
pytestmark = pytest.mark.integration
|
||
|
||
|
||
_TEST_TEXT = "Договор поставки. " * 50
|
||
|
||
|
||
class FakePublisher:
|
||
"""Captures the message published by ``send_to_analysis``."""
|
||
|
||
def __init__(self) -> None:
|
||
self.messages: list[tuple[Any, str]] = []
|
||
|
||
async def publish(self, message: Any, routing_key: str) -> None:
|
||
self.messages.append((message, routing_key))
|
||
|
||
|
||
class StubProvider:
|
||
"""LLMProvider returning one fixed finding."""
|
||
|
||
async def analyze(self, text: str, *, checklist: str) -> AnalysisResult:
|
||
return AnalysisResult(
|
||
findings=[
|
||
Finding(
|
||
checklist_id="penalties",
|
||
severity="high",
|
||
quote="штраф 0,5%",
|
||
section_ref="п. 6.3",
|
||
risk="высокая неустойка",
|
||
recommendation="ограничить",
|
||
)
|
||
],
|
||
model_used="stub-model",
|
||
prompt_tokens=10,
|
||
eval_tokens=20,
|
||
latency_sec=0.1,
|
||
)
|
||
|
||
async def extract_prescreen(self, text: str, *, max_chars: int | None = None) -> dict[str, Any]:
|
||
raise NotImplementedError
|
||
|
||
async def aclose(self) -> None:
|
||
pass
|
||
|
||
|
||
def _storage(infra: dict[str, str]) -> MinioStorage:
|
||
return MinioStorage.from_endpoint_url(
|
||
endpoint_url=infra["s3_endpoint_url"],
|
||
access_key=infra["s3_access_key"],
|
||
secret_key=infra["s3_secret_key"],
|
||
bucket=infra["s3_bucket"],
|
||
)
|
||
|
||
|
||
def _session_factory() -> Any:
|
||
return create_session_factory()
|
||
|
||
|
||
async def _seed_user(session: Any, *, telegram_id: int, credits: int = 5) -> uuid.UUID:
|
||
result = await session.execute(
|
||
text(
|
||
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, :c) "
|
||
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = :c "
|
||
"RETURNING id"
|
||
),
|
||
{"t": telegram_id, "c": credits},
|
||
)
|
||
return result.scalar_one()
|
||
|
||
|
||
async def _seed_manual_review_doc(
|
||
infra: dict[str, str],
|
||
*,
|
||
telegram_id: int | None = None,
|
||
document_id: uuid.UUID,
|
||
confidence: float = 0.5,
|
||
total_amount: Decimal = Decimal("100000.00"),
|
||
) -> tuple[uuid.UUID, str, uuid.UUID]:
|
||
"""Create user + extracted text in MinIO + manual_review document + prescreen result."""
|
||
if telegram_id is None:
|
||
telegram_id = int(uuid.uuid4().int % 1_000_000_000)
|
||
sess = _session_factory()
|
||
store = _storage(infra)
|
||
async with sess() as session:
|
||
user_id = await _seed_user(session, telegram_id=telegram_id)
|
||
await session.commit()
|
||
|
||
ext_key = extracted_key(str(user_id), str(document_id))
|
||
await store.put(ext_key, _TEST_TEXT.encode("utf-8"), content_type="text/plain; charset=utf-8")
|
||
|
||
correlation_id = uuid.uuid4()
|
||
async with sess() as session:
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO documents "
|
||
"(id, user_id, s3_key, extracted_s3_key, filename, mime, bytes, status, stage) "
|
||
"VALUES (:id, :uid, :s3, :ext, 'contract.pdf', 'application/pdf', :bytes, "
|
||
" 'manual_review', 'manual_review')"
|
||
),
|
||
{
|
||
"id": document_id,
|
||
"uid": user_id,
|
||
"s3": f"users/{user_id}/docs/{document_id}.pdf",
|
||
"ext": ext_key,
|
||
"bytes": len(_TEST_TEXT),
|
||
},
|
||
)
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO prescreen_results "
|
||
"(id, document_id, correlation_id, contract_type, party_a, party_b, "
|
||
" total_amount, currency, start_date, end_date, has_penalty_clause, "
|
||
" has_termination_clause, has_arbitration, confidence_score, routing_decision, "
|
||
" prescreened_at, auto_summary, auto_findings, extractor_version) "
|
||
"VALUES (:id, :d, :cid, 'supply', 'ООО Продавец', 'ООО Покупатель', "
|
||
" :ta, 'RUB', '2025-01-01', '2025-12-31', TRUE, TRUE, FALSE, "
|
||
" :cs, 'manual_review', now(), 'summary', '[{\"note\":\"ok\"}]', 'heuristic-v2')"
|
||
),
|
||
{
|
||
"id": uuid.uuid4(),
|
||
"d": document_id,
|
||
"cid": correlation_id,
|
||
"ta": total_amount,
|
||
"cs": confidence,
|
||
},
|
||
)
|
||
await session.commit()
|
||
return user_id, ext_key, correlation_id
|
||
|
||
|
||
async def _doc_state(session: Any, document_id: uuid.UUID) -> tuple[str, str, bool]:
|
||
result = await session.execute(
|
||
text("SELECT status, stage, refunded FROM documents WHERE id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
return result.one()
|
||
|
||
|
||
async def test_reject_credits_refunds_once(infra: dict[str, str]) -> None:
|
||
sess = _session_factory()
|
||
document_id = uuid.uuid4()
|
||
user_id, _ext_key, _correlation_id = await _seed_manual_review_doc(
|
||
infra, document_id=document_id
|
||
)
|
||
|
||
async with sess() as session:
|
||
await CreditsRepository(session).reserve(user_id, document_id=document_id)
|
||
await session.commit()
|
||
|
||
async with sess() as session:
|
||
await reject(session, document_id=document_id, reason="Некорректный договор")
|
||
|
||
async with sess() as session:
|
||
status, stage, refunded = await _doc_state(session, document_id)
|
||
assert status == "failed"
|
||
assert stage == "review_rejected"
|
||
assert refunded is True
|
||
|
||
balance = await CreditsRepository(session).get_balance(user_id)
|
||
assert balance == 5
|
||
|
||
events = await session.execute(
|
||
text("SELECT kind, delta FROM credit_events WHERE user_id = :u AND document_id = :d"),
|
||
{"u": user_id, "d": document_id},
|
||
)
|
||
rows = events.all()
|
||
assert any(kind == "reserve" for kind, _ in rows)
|
||
assert any(kind == "refund_auto" and delta == 1 for kind, delta in rows)
|
||
|
||
# Second reject is idempotent via status, not compensation again.
|
||
async with sess() as session:
|
||
with pytest.raises(ReviewActionError):
|
||
await reject(session, document_id=document_id, reason="")
|
||
|
||
|
||
async def test_reject_quota_restores_slot_no_credit_refund(infra: dict[str, str]) -> None:
|
||
sess = _session_factory()
|
||
document_id = uuid.uuid4()
|
||
user_id, _ext_key, _correlation_id = await _seed_manual_review_doc(
|
||
infra, document_id=document_id
|
||
)
|
||
|
||
subscription_id = uuid.uuid4()
|
||
async with sess() as session:
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO subscriptions "
|
||
"(id, user_id, plan_code, status, current_period_start, current_period_end) "
|
||
"VALUES (:id, :u, 'lite', 'active', now() - interval '1 day', now() + interval '30 days')"
|
||
),
|
||
{"id": subscription_id, "u": user_id},
|
||
)
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO quota_usage (user_id, subscription_id, document_id) "
|
||
"VALUES (:u, :s, :d)"
|
||
),
|
||
{"u": user_id, "s": subscription_id, "d": document_id},
|
||
)
|
||
await session.commit()
|
||
|
||
async with sess() as session:
|
||
await reject(session, document_id=document_id, reason="Неподходящий формат")
|
||
|
||
async with sess() as session:
|
||
status, stage, refunded = await _doc_state(session, document_id)
|
||
assert status == "failed"
|
||
assert stage == "review_rejected"
|
||
assert refunded is True
|
||
|
||
quota = await session.execute(
|
||
text("SELECT id FROM quota_usage WHERE document_id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
assert quota.first() is None
|
||
|
||
events = await session.execute(
|
||
text("SELECT kind, delta FROM credit_events WHERE user_id = :u AND document_id = :d"),
|
||
{"u": user_id, "d": document_id},
|
||
)
|
||
assert not any(kind == "refund_auto" for kind, _ in events.all())
|
||
|
||
balance = await CreditsRepository(session).get_balance(user_id)
|
||
assert balance == 5 # default seeded credits, untouched
|
||
|
||
|
||
async def test_complete_writes_report_and_done(infra: dict[str, str]) -> None:
|
||
sess = _session_factory()
|
||
document_id = uuid.uuid4()
|
||
user_id, _ext_key, _correlation_id = await _seed_manual_review_doc(
|
||
infra, document_id=document_id
|
||
)
|
||
|
||
async with sess() as session:
|
||
await complete(session, document_id=document_id, note="Договор типовой, рисков нет.")
|
||
|
||
async with sess() as session:
|
||
status, stage, _refunded = await _doc_state(session, document_id)
|
||
assert status == "done"
|
||
assert stage == "review_completed"
|
||
|
||
report = await ReportRepository(session).get_by_document_id(document_id)
|
||
assert report is not None
|
||
assert report.model_used == "admin-complete"
|
||
assert "Договор типовой, рисков нет." in report.markdown
|
||
assert report.prescreen_result_id is not None
|
||
assert report.prescreen_meta is not None
|
||
|
||
|
||
async def test_send_to_analysis_publishes_message_and_handler_finishes(
|
||
infra: dict[str, str],
|
||
) -> None:
|
||
sess = _session_factory()
|
||
document_id = uuid.uuid4()
|
||
user_id, ext_key, _correlation_id = await _seed_manual_review_doc(
|
||
infra, document_id=document_id
|
||
)
|
||
|
||
fake_publisher = FakePublisher()
|
||
|
||
async with sess() as session:
|
||
await send_to_analysis(session, fake_publisher, document_id=document_id)
|
||
|
||
assert len(fake_publisher.messages) == 1
|
||
msg, routing_key = fake_publisher.messages[0]
|
||
assert isinstance(msg, AnalyzeRequested)
|
||
assert msg.document_id == document_id
|
||
assert msg.user_id == user_id
|
||
assert msg.extracted_s3_key == ext_key
|
||
assert msg.attempt == 0
|
||
assert msg.prescreen_meta is not None
|
||
assert msg.prescreen_result_id is not None
|
||
assert routing_key == RK_ANALYZE
|
||
|
||
async with sess() as session:
|
||
status, stage, _refunded = await _doc_state(session, document_id)
|
||
assert status == "analyzing"
|
||
assert stage == "queued_analyze"
|
||
|
||
# Drive the analyze worker with the same message + stub provider.
|
||
analyze_handler = AnalyzeHandler(session_factory=sess, provider=StubProvider())
|
||
try:
|
||
await analyze_handler.handle(msg)
|
||
finally:
|
||
await analyze_handler.aclose()
|
||
|
||
async with sess() as session:
|
||
status, stage, _refunded = await _doc_state(session, document_id)
|
||
assert status == "done"
|
||
assert stage == "done"
|
||
|
||
report = await ReportRepository(session).get_by_document_id(document_id)
|
||
assert report is not None
|
||
assert len(report.content_json.get("findings", [])) == 1
|
||
|
||
|
||
async def test_send_to_analysis_rolls_back_status_on_publish_failure(
|
||
infra: dict[str, str],
|
||
) -> None:
|
||
sess = _session_factory()
|
||
document_id = uuid.uuid4()
|
||
_user_id, _ext_key, _correlation_id = await _seed_manual_review_doc(
|
||
infra, document_id=document_id
|
||
)
|
||
|
||
class FailingPublisher:
|
||
async def publish(self, message: Any, routing_key: str) -> None:
|
||
raise RuntimeError("broker down")
|
||
|
||
async with sess() as session:
|
||
with pytest.raises(RuntimeError):
|
||
await send_to_analysis(session, FailingPublisher(), document_id=document_id)
|
||
|
||
async with sess() as session:
|
||
status, stage, _refunded = await _doc_state(session, document_id)
|
||
assert status == "manual_review"
|
||
assert stage == "manual_review"
|
||
|
||
|
||
async def test_action_requires_manual_review_status(infra: dict[str, str]) -> None:
|
||
sess = _session_factory()
|
||
document_id = uuid.uuid4()
|
||
_user_id, _ext_key, _correlation_id = await _seed_manual_review_doc(
|
||
infra, document_id=document_id
|
||
)
|
||
|
||
async with sess() as session:
|
||
await DocumentRepository(session).update_status(document_id, status="done")
|
||
await session.commit()
|
||
|
||
async with sess() as session:
|
||
with pytest.raises(ReviewActionError):
|
||
await complete(session, document_id=document_id, note="note")
|