"""Integration tests for worker-prescreen (hybrid extraction). Uses the same Docker Compose infra as test_analyze_worker: real Postgres, RabbitMQ, MinIO; the LLM is stubbed at the provider port. Scenarios (plan Phase 4): 1. Heuristic-only (fallback disabled by default): prescreen_results row persisted with extractor_version='heuristic-v2', routing unchanged. 2. Fallback enabled + low-confidence text + stub provider: merged meta persisted with extractor_version='hybrid-llm-v1'. 3. Fallback failure: heuristic result kept, error recorded in auto_findings. """ from __future__ import annotations import json import uuid from typing import TYPE_CHECKING, Any import pytest from sqlalchemy import text from contract_check.core.db.session import create_session_factory from contract_check.core.mq.messages import PrescreenRequested from contract_check.core.s3 import extracted_key from contract_check.core.s3.minio_storage import MinioStorage from contract_check.worker_prescreen.extractor_heuristic import HeuristicExtractor from contract_check.worker_prescreen.extractor_hybrid import HybridMetaExtractor from contract_check.worker_prescreen.handler import PrescreenHandler if TYPE_CHECKING: from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker pytestmark = pytest.mark.integration FULL_CONTRACT = """\ ДОГОВОР ПОСТАВКИ № 42 г. Москва Общество с ограниченной ответственностью «Продавец», именуемое в дальнейшем «Поставщик», с одной стороны, и Общество с ограниченной ответственностью «Покупатель», с другой стороны, заключили настоящий договор. 1. Предмет договора Поставщик обязуется передать в собственность Покупателю товар. 2. Цена и порядок расчётов 2.1. Общая стоимость товара составляет 1 250 000 рублей. 3. Срок действия договора 3.1. Договор вступает в силу с 01.09.2025 и действует по 31.08.2026. 4. Ответственность сторон 4.1. За просрочку поставки Поставщик уплачивает неустойку. 5. Порядок разрешения споров 5.1. Споры подлежат рассмотрению в Арбитражном суде г. Москвы. """ LOW_CONF_TEXT = "Скан документа без распознанных реквизитов." # The compose stack may have live workers consuming analyze.q; publish test # messages to an unrouted key so nothing downstream races our assertions. _DEAD_ROUTING_KEY = "it-test-unrouted" class StubProvider: """LLMProvider stub returning a fixed prescreen dict (or raising).""" def __init__(self, payload: dict[str, Any] | Exception) -> None: self._payload = payload async def analyze(self, text: str, *, checklist: str, extra_context: str = "") -> Any: raise NotImplementedError async def extract_prescreen(self, text: str) -> dict[str, Any]: if isinstance(self._payload, Exception): raise self._payload return self._payload async def aclose(self) -> None: pass LLM_DICT: dict[str, Any] = { "contract_type": "services", "party_a": "ИП Петров Петр Петрович", "party_b": "ООО «Клиент»", "total_amount": 50000.0, "currency": "RUB", "has_penalty_clause": True, } 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="contract-check-docs", ) async def _seed_document( infra: dict[str, str], *, telegram_id: int, document_id: uuid.UUID, contract_text: str ) -> tuple[uuid.UUID, str]: """Insert user + post-extraction document; return (user_id, text_s3_key).""" sess = create_session_factory() store = _storage(infra) async with sess() as session: result = await session.execute( text( "INSERT INTO users (telegram_id, credits_left) VALUES (:t, 5) " "ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 5 " "RETURNING id" ), {"t": telegram_id}, ) user_id = result.scalar_one() await session.commit() ext_key = extracted_key(str(user_id), str(document_id)) 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, 'contract.pdf', 'application/pdf', :bytes, " " 'prescreening', 'extracted')" ), { "id": document_id, "uid": user_id, "s3": f"users/{user_id}/docs/{document_id}.pdf", "ext": ext_key, "bytes": len(contract_text), }, ) await session.commit() return user_id, ext_key def _requested( payload_doc: uuid.UUID, user_id: uuid.UUID, ext_key: str, text: str ) -> PrescreenRequested: return PrescreenRequested( correlation_id=uuid.uuid4(), document_id=payload_doc, user_id=user_id, text_s3_key=ext_key, filename="contract.pdf", char_count=len(text), ) async def _prescreen_row(sess: async_sessionmaker[AsyncSession], document_id: uuid.UUID) -> Any: async with sess() as session: result = await session.execute( text( "SELECT contract_type, party_a, total_amount, currency, confidence_score, " " routing_decision, extractor_version, auto_findings " "FROM prescreen_results WHERE document_id = :d" ), {"d": document_id}, ) return result.one() async def test_prescreen_heuristic_persists_version_and_routes( infra: dict[str, str], ) -> None: sess = create_session_factory() document_id = uuid.uuid4() user_id, ext_key = await _seed_document( infra, telegram_id=777_111_222, document_id=document_id, contract_text=FULL_CONTRACT ) handler = PrescreenHandler(session_factory=sess, publish_routing_key=_DEAD_ROUTING_KEY) try: await handler.handle(_requested(document_id, user_id, ext_key, FULL_CONTRACT)) finally: await handler.aclose() row = await _prescreen_row(sess, document_id) (ctype, party_a, amount, currency, confidence, decision, version, findings) = row assert ctype == "supply" assert "Продавец" in party_a assert amount == 1250000.0 assert currency == "RUB" assert confidence == 1.0 assert decision == "deep_analysis" # high-value + penalty + arbitration assert version == "heuristic-v2" # fallback disabled by default assert findings == [] async with sess() as session: doc = await session.execute( text("SELECT status, stage FROM documents WHERE id = :d"), {"d": document_id} ) status, stage = doc.one() assert status == "analyzing" assert stage == "queued_analyze" async def test_prescreen_llm_fallback_merged_and_persisted( infra: dict[str, str], ) -> None: sess = create_session_factory() document_id = uuid.uuid4() user_id, ext_key = await _seed_document( infra, telegram_id=777_333_444, document_id=document_id, contract_text=LOW_CONF_TEXT ) handler = PrescreenHandler( session_factory=sess, publish_routing_key=_DEAD_ROUTING_KEY, extractor=HybridMetaExtractor( provider=StubProvider(LLM_DICT), fallback_enabled=True, fallback_threshold=0.99, heuristic=HeuristicExtractor(), ), ) await handler.handle(_requested(document_id, user_id, ext_key, LOW_CONF_TEXT)) row = await _prescreen_row(sess, document_id) (ctype, party_a, amount, currency, confidence, decision, version, findings) = row assert ctype == "services" # recovered by the LLM stage assert party_a == "ИП Петров Петр Петрович" assert amount == 50000.0 assert currency == "RUB" assert confidence >= 0.75 assert decision == "deep_analysis" # has_penalty_clause OR-merged to True assert version == "hybrid-llm-v1" assert findings == [] async def test_prescreen_llm_failure_records_error_not_fatal( infra: dict[str, str], ) -> None: sess = create_session_factory() document_id = uuid.uuid4() user_id, ext_key = await _seed_document( infra, telegram_id=777_555_666, document_id=document_id, contract_text=LOW_CONF_TEXT ) handler = PrescreenHandler( session_factory=sess, publish_routing_key=_DEAD_ROUTING_KEY, extractor=HybridMetaExtractor( provider=StubProvider(RuntimeError("quota exceeded")), fallback_enabled=True, fallback_threshold=0.99, heuristic=HeuristicExtractor(), ), ) # Must not raise: the heuristic result survives the LLM failure. await handler.handle(_requested(document_id, user_id, ext_key, LOW_CONF_TEXT)) row = await _prescreen_row(sess, document_id) (ctype, _party_a, _amount, _cur, confidence, decision, version, findings) = row assert ctype is None assert confidence < 0.75 assert decision == "manual_review" # routing semantics unchanged assert version == "heuristic-v2" parsed = findings if isinstance(findings, list) else json.loads(findings) assert parsed and "llm_fallback_error" in parsed[0] assert "quota exceeded" in parsed[0]["llm_fallback_error"] async with sess() as session: doc = await session.execute( text("SELECT status, stage FROM documents WHERE id = :d"), {"d": document_id} ) status, stage = doc.one() assert status == "manual_review" assert stage == "manual_review" async def test_prescreen_redelivery_for_manual_review_is_no_op( infra: dict[str, str], ) -> None: sess = create_session_factory() document_id = uuid.uuid4() user_id, ext_key = await _seed_document( infra, telegram_id=777_555_667, document_id=document_id, contract_text=LOW_CONF_TEXT ) handler = PrescreenHandler( session_factory=sess, publish_routing_key=_DEAD_ROUTING_KEY, extractor=HybridMetaExtractor( provider=StubProvider(RuntimeError("quota exceeded")), fallback_enabled=True, fallback_threshold=0.99, heuristic=HeuristicExtractor(), ), ) await handler.handle(_requested(document_id, user_id, ext_key, LOW_CONF_TEXT)) # Force-update a field so we can detect a re-run. async with sess() as session: await session.execute( text("UPDATE prescreen_results SET confidence_score = 0.123 WHERE document_id = :d"), {"d": document_id}, ) await session.commit() # Redeliver the same prescreen message: it should be a no-op. await handler.handle(_requested(document_id, user_id, ext_key, LOW_CONF_TEXT)) async with sess() as session: result = await session.execute( text("SELECT confidence_score FROM prescreen_results WHERE document_id = :d"), {"d": document_id}, ) assert float(result.scalar_one()) == pytest.approx(0.123)