102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
"""Unit tests for worker-extract handler.
|
|
|
|
Covers the `PRESCREEN_ENABLED` toggle: when true it publishes
|
|
`PrescreenRequested` to `prescreen.q` and marks `prescreening`; when false it
|
|
bypasses the prescreen stage, publishes `AnalyzeRequested` to `analyze.q`,
|
|
and marks `analyzing`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from contract_check.core.extraction import ExtractedDocument
|
|
from contract_check.core.mq.messages import DocumentExtracted, DocumentUploaded, PrescreenRequested
|
|
from contract_check.worker_extract.handler import ExtractHandler
|
|
|
|
|
|
def _fake_session_factory() -> Any:
|
|
session = MagicMock()
|
|
session.execute = AsyncMock()
|
|
session.commit = AsyncMock()
|
|
session.__aenter__ = AsyncMock(return_value=session)
|
|
session.__aexit__ = AsyncMock(return_value=False)
|
|
return MagicMock(return_value=session), session
|
|
|
|
|
|
def _uploaded() -> DocumentUploaded:
|
|
return DocumentUploaded(
|
|
correlation_id=uuid.uuid4(),
|
|
document_id=uuid.uuid4(),
|
|
user_id=uuid.uuid4(),
|
|
s3_key="users/u/docs/d.pdf",
|
|
filename="contract.pdf",
|
|
mime="application/pdf",
|
|
)
|
|
|
|
|
|
def _extracted() -> ExtractedDocument:
|
|
return ExtractedDocument(
|
|
markdown="Договор поставки. Сумма 1000 рублей.",
|
|
metadata={"format": "pdf", "is_scan": False, "has_tables": False},
|
|
is_structured=True,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def handler(monkeypatch: pytest.MonkeyPatch) -> tuple[ExtractHandler, Any, Any]:
|
|
session_factory, session = _fake_session_factory()
|
|
h = ExtractHandler(session_factory=session_factory)
|
|
|
|
# Stub S3 put and the extraction core.
|
|
h._storage = MagicMock(put=AsyncMock())
|
|
monkeypatch.setattr(h, "_extract_text", AsyncMock(return_value=(_extracted().markdown, False)))
|
|
|
|
# Stub publisher.
|
|
publisher = MagicMock(publish=AsyncMock())
|
|
h._publisher = publisher
|
|
|
|
# DB status is non-terminal.
|
|
monkeypatch.setattr(h, "_db_status", AsyncMock(return_value="queued"))
|
|
|
|
return h, session, publisher
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prescreen_enabled_publishes_to_prescreen(
|
|
handler: tuple[ExtractHandler, Any, Any],
|
|
) -> None:
|
|
h, session, publisher = handler
|
|
h._settings.prescreen_enabled = True
|
|
|
|
await h.handle(_uploaded())
|
|
|
|
assert publisher.publish.call_count == 1
|
|
msg, *_ = publisher.publish.call_args.args
|
|
assert isinstance(msg, PrescreenRequested)
|
|
assert publisher.publish.call_args.kwargs["routing_key"] == "prescreen"
|
|
|
|
status_params = session.execute.call_args_list[-2].args[1]
|
|
assert status_params["status"] == "prescreening"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prescreen_disabled_bypasses_to_analyze(
|
|
handler: tuple[ExtractHandler, Any, Any],
|
|
) -> None:
|
|
h, session, publisher = handler
|
|
h._settings.prescreen_enabled = False
|
|
|
|
await h.handle(_uploaded())
|
|
|
|
assert publisher.publish.call_count == 1
|
|
msg, *_ = publisher.publish.call_args.args
|
|
assert isinstance(msg, DocumentExtracted)
|
|
assert publisher.publish.call_args.kwargs["routing_key"] == "analyze"
|
|
|
|
status_params = session.execute.call_args_list[-2].args[1]
|
|
assert status_params["status"] == "analyzing"
|