"""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