DealDocumentScreening/tests/integration/test_prescreen_worker.py

"""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=infra["s3_bucket"],
    )


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)