265 lines
8.7 KiB
Python
265 lines
8.7 KiB
Python
"""Integration tests for worker-analyze.
|
|
|
|
Uses the same Docker Compose infra as test_upload_pipeline / test_extract_worker:
|
|
real Postgres, RabbitMQ, MinIO. The LLM is respx-mocked at a fake Ollama host.
|
|
|
|
Scenarios:
|
|
1. Happy path: DocumentExtracted -> report saved (JSONB+markdown), status=done,
|
|
disclaimer present, jobs row done.
|
|
2. Terminal failure: LLM quota error -> on_terminal_failure refunds + DLQ state.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
from typing import TYPE_CHECKING
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
from sqlalchemy import text
|
|
|
|
from contract_check.core.db.session import create_session_factory
|
|
from contract_check.core.llm.ollama_cloud import LLMQuotaError, OllamaCloudProvider
|
|
from contract_check.core.mq.messages import AnalyzeRequested
|
|
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
|
|
|
|
if TYPE_CHECKING:
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
HOST = "https://ollama.test"
|
|
URL = f"{HOST}/api/chat"
|
|
|
|
_VALID_FINDINGS = {
|
|
"findings": [
|
|
{
|
|
"checklist_id": "penalties",
|
|
"severity": "high",
|
|
"quote": "Штраф 0,5% за каждый день просрочки",
|
|
"section_ref": "п. 6.3",
|
|
"risk": "Высокая неустойка",
|
|
"recommendation": "Ограничить cap",
|
|
}
|
|
]
|
|
}
|
|
|
|
|
|
def _chat_body(content: str, model: str = "qwen2.5:14b") -> dict[str, object]:
|
|
return {
|
|
"model": model,
|
|
"message": {"content": content},
|
|
"prompt_eval_count": 10,
|
|
"eval_count": 20,
|
|
}
|
|
|
|
|
|
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(infra: dict[str, str]) -> async_sessionmaker[AsyncSession]:
|
|
return create_session_factory()
|
|
|
|
|
|
def _provider() -> OllamaCloudProvider:
|
|
return OllamaCloudProvider(
|
|
host=HOST,
|
|
api_key="key",
|
|
model="qwen2.5:14b",
|
|
fallback_model=None,
|
|
max_concurrency=1,
|
|
chunk_size=10000,
|
|
)
|
|
|
|
|
|
async def _seed_document(
|
|
infra: dict[str, str], *, telegram_id: int, credits: int, document_id: uuid.UUID
|
|
) -> tuple[uuid.UUID, int, str, str]:
|
|
"""Insert user + a post-extraction document; return (user_id, text, ext_key, s3_key)."""
|
|
sess = _session_factory(infra)
|
|
store = _storage(infra)
|
|
async with sess() as session:
|
|
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, credits_left"
|
|
),
|
|
{"t": telegram_id, "c": credits},
|
|
)
|
|
row = result.first()
|
|
assert row is not None
|
|
user_id, initial_credits = row
|
|
await session.commit()
|
|
|
|
ext_key = extracted_key(str(user_id), str(document_id))
|
|
s3_key = f"users/{user_id}/docs/{document_id}.pdf"
|
|
contract_text = "Договор поставки. " * 50
|
|
await store.put(
|
|
ext_key, contract_text.encode("utf-8"), content_type="text/plain; charset=utf-8"
|
|
)
|
|
|
|
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, :fn, :mime, :bytes, "
|
|
" 'analyzing', 'queued_analyze')"
|
|
),
|
|
{
|
|
"id": document_id,
|
|
"uid": user_id,
|
|
"s3": s3_key,
|
|
"ext": ext_key,
|
|
"fn": "contract.pdf",
|
|
"mime": "application/pdf",
|
|
"bytes": len(contract_text),
|
|
},
|
|
)
|
|
await session.execute(
|
|
text(
|
|
"INSERT INTO jobs (document_id, correlation_id, queue, status) "
|
|
"VALUES (:did, :cid, 'analyze', 'pending')"
|
|
),
|
|
{"did": document_id, "cid": uuid.uuid4()},
|
|
)
|
|
await session.commit()
|
|
|
|
return user_id, initial_credits, ext_key, contract_text
|
|
|
|
|
|
async def test_analyze_worker_saves_report_and_marks_done(
|
|
infra: dict[str, str],
|
|
) -> None:
|
|
sess = _session_factory(infra)
|
|
document_id = uuid.uuid4()
|
|
correlation_id = uuid.uuid4()
|
|
user_id, _initial, ext_key, contract_text = await _seed_document(
|
|
infra, telegram_id=222_333_444, credits=5, document_id=document_id
|
|
)
|
|
|
|
async with _provider() as provider:
|
|
with respx.mock(base_url=HOST) as mock:
|
|
mock.post(URL).mock(
|
|
return_value=httpx.Response(200, json=_chat_body(json.dumps(_VALID_FINDINGS)))
|
|
)
|
|
await AnalyzeHandler(session_factory=sess, provider=provider).handle(
|
|
AnalyzeRequested(
|
|
correlation_id=correlation_id,
|
|
document_id=document_id,
|
|
user_id=user_id,
|
|
extracted_s3_key=ext_key,
|
|
char_count=len(contract_text),
|
|
ocr_used=False,
|
|
)
|
|
)
|
|
|
|
async with sess() as session:
|
|
report = await session.execute(
|
|
text(
|
|
"SELECT markdown, model_used, prompt_tokens, eval_tokens "
|
|
"FROM reports WHERE document_id = :d"
|
|
),
|
|
{"d": document_id},
|
|
)
|
|
row = report.first()
|
|
assert row is not None
|
|
markdown, model_used, prompt_tokens, eval_tokens = row
|
|
assert "Дисклеймер" in markdown
|
|
assert "Штраф 0,5%" in markdown
|
|
assert model_used == "qwen2.5:14b"
|
|
assert prompt_tokens == 10
|
|
assert eval_tokens == 20
|
|
|
|
doc = await session.execute(
|
|
text("SELECT status, stage FROM documents WHERE id = :d"),
|
|
{"d": document_id},
|
|
)
|
|
status, stage = doc.one()
|
|
assert status == "done"
|
|
assert stage == "done"
|
|
|
|
job = await session.execute(
|
|
text("SELECT status FROM jobs WHERE document_id = :d AND queue = 'analyze'"),
|
|
{"d": document_id},
|
|
)
|
|
assert job.scalar_one() == "done"
|
|
|
|
|
|
async def test_analyze_worker_terminal_failure_refunds_and_dlqs(
|
|
infra: dict[str, str],
|
|
) -> None:
|
|
sess = _session_factory(infra)
|
|
document_id = uuid.uuid4()
|
|
correlation_id = uuid.uuid4()
|
|
user_id, initial_credits, ext_key, contract_text = await _seed_document(
|
|
infra, telegram_id=555_666_777, credits=3, document_id=document_id
|
|
)
|
|
|
|
async with _provider() as provider:
|
|
handler = AnalyzeHandler(session_factory=sess, provider=provider)
|
|
with respx.mock(base_url=HOST) as mock:
|
|
mock.post(URL).mock(return_value=httpx.Response(429, json={"error": "quota"}))
|
|
with pytest.raises(LLMQuotaError):
|
|
await handler.handle(
|
|
AnalyzeRequested(
|
|
correlation_id=correlation_id,
|
|
document_id=document_id,
|
|
user_id=user_id,
|
|
extracted_s3_key=ext_key,
|
|
char_count=len(contract_text),
|
|
ocr_used=False,
|
|
)
|
|
)
|
|
|
|
await handler.on_terminal_failure(
|
|
AnalyzeRequested(
|
|
correlation_id=correlation_id,
|
|
document_id=document_id,
|
|
user_id=user_id,
|
|
extracted_s3_key=ext_key,
|
|
char_count=len(contract_text),
|
|
ocr_used=False,
|
|
),
|
|
"llm_quota",
|
|
"test terminal failure",
|
|
)
|
|
|
|
async with sess() as session:
|
|
doc = await session.execute(
|
|
text("SELECT status, refunded FROM documents WHERE id = :d"),
|
|
{"d": document_id},
|
|
)
|
|
status, refunded = doc.one()
|
|
assert status == "failed"
|
|
assert refunded is True
|
|
|
|
job = await session.execute(
|
|
text(
|
|
"SELECT status, dlq, last_failure_class FROM jobs "
|
|
"WHERE document_id = :d AND queue = 'analyze'"
|
|
),
|
|
{"d": document_id},
|
|
)
|
|
j_status, j_dlq, j_class = job.one()
|
|
assert j_status == "dlq"
|
|
assert j_dlq is True
|
|
assert j_class == "llm_quota"
|
|
|
|
credits = await session.execute(
|
|
text("SELECT credits_left FROM users WHERE id = :u"),
|
|
{"u": user_id},
|
|
)
|
|
assert credits.scalar_one() == initial_credits + 1
|