377 lines
13 KiB
Python
377 lines
13 KiB
Python
"""Integration tests for worker-extract.
|
||
|
||
Uses the same Docker Compose infra as test_upload_pipeline:
|
||
- real Postgres, RabbitMQ, MinIO
|
||
- alembic migrations + seeded service token (via conftest.py infra fixture)
|
||
|
||
Scenarios:
|
||
1. Happy path: DocumentUploaded → extracted text on analyze.q + status=analyzing.
|
||
2. DOCX extraction also works end-to-end.
|
||
3. Terminal failure path: bad PDF → extraction_failed → refund + document failed.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import TYPE_CHECKING
|
||
|
||
import aio_pika
|
||
import pytest
|
||
|
||
from contract_check.core.db.session import create_session_factory
|
||
from contract_check.core.mq.messages import DocumentUploaded
|
||
from contract_check.core.s3 import original_key
|
||
from contract_check.core.s3.minio_storage import MinioStorage
|
||
from contract_check.worker_extract.handler import ExtractHandler
|
||
|
||
if TYPE_CHECKING:
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||
|
||
pytestmark = pytest.mark.integration
|
||
|
||
|
||
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",
|
||
)
|
||
|
||
|
||
def session_factory(infra: dict[str, str]) -> async_sessionmaker[AsyncSession]:
|
||
return create_session_factory()
|
||
|
||
|
||
@pytest.fixture
|
||
def pdf_bytes(tmp_path: Path) -> bytes:
|
||
import pymupdf
|
||
|
||
doc = pymupdf.open()
|
||
page = doc.new_page()
|
||
long_text = (
|
||
"Договор. Стороны обязуются выполнять условия. "
|
||
"Сторона А обязуется передать товар. "
|
||
"Сторона Б обязуется оплатить товар в течение десяти банковских дней. "
|
||
"Ответственность сторон ограничена суммой договора. "
|
||
"Споры подлежат рассмотрению в арбитражном суде города Москвы."
|
||
)
|
||
page.insert_htmlbox(
|
||
page.rect,
|
||
f'<p style="font-family:DejaVu Sans;font-size:14px">{long_text}</p>',
|
||
)
|
||
path = tmp_path / "contract.pdf"
|
||
doc.save(str(path))
|
||
doc.close()
|
||
return path.read_bytes()
|
||
|
||
|
||
@pytest.fixture
|
||
def docx_bytes(tmp_path: Path) -> bytes:
|
||
from docx import Document
|
||
|
||
doc = Document()
|
||
long_text = (
|
||
"Договор. Стороны обязуются выполнять условия. "
|
||
"Сторона А обязуется передать товар. "
|
||
"Сторона Б обязуется оплатить товар в течение десяти банковских дней. "
|
||
"Ответственность сторон ограничена суммой договора. "
|
||
"Споры подлежат рассмотрению в арбитражном суде города Москвы."
|
||
)
|
||
doc.add_paragraph(long_text)
|
||
path = tmp_path / "contract.docx"
|
||
doc.save(str(path))
|
||
return path.read_bytes()
|
||
|
||
|
||
async def test_extract_worker_pdf_uploads_text_and_publishes_analyze(
|
||
infra: dict[str, str],
|
||
pdf_bytes: bytes,
|
||
) -> None:
|
||
from sqlalchemy import text
|
||
|
||
telegram_id = 111_222_333
|
||
user_id: uuid.UUID | None = None
|
||
sess_factory = session_factory(infra)
|
||
store = storage(infra)
|
||
async with sess_factory() 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()
|
||
assert user_id is not None
|
||
|
||
document_id = uuid.uuid4()
|
||
correlation_id = uuid.uuid4()
|
||
s3_key = original_key(str(user_id), str(document_id), ".pdf")
|
||
await store.put(s3_key, pdf_bytes, content_type="application/pdf")
|
||
|
||
async with sess_factory() as session:
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO documents (id, user_id, s3_key, filename, mime, bytes, status) "
|
||
"VALUES (:id, :uid, :s3, :fn, :mime, :bytes, 'queued')"
|
||
),
|
||
{
|
||
"id": document_id,
|
||
"uid": user_id,
|
||
"s3": s3_key,
|
||
"fn": "contract.pdf",
|
||
"mime": "application/pdf",
|
||
"bytes": len(pdf_bytes),
|
||
},
|
||
)
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO jobs (document_id, correlation_id, queue, status) "
|
||
"VALUES (:did, :cid, 'extract', 'pending')"
|
||
),
|
||
{"did": document_id, "cid": correlation_id},
|
||
)
|
||
await session.commit()
|
||
|
||
handler = ExtractHandler(session_factory=sess_factory)
|
||
await handler.handle(
|
||
DocumentUploaded(
|
||
correlation_id=correlation_id,
|
||
document_id=document_id,
|
||
user_id=user_id,
|
||
s3_key=s3_key,
|
||
filename="contract.pdf",
|
||
mime="application/pdf",
|
||
)
|
||
)
|
||
|
||
extracted_key_path = f"users/{user_id}/docs/{document_id}.txt"
|
||
text_data = await store.get(extracted_key_path)
|
||
assert "Договор" in text_data.decode("utf-8")
|
||
|
||
connection = await aio_pika.connect_robust(infra["rabbitmq_url"])
|
||
try:
|
||
channel = await connection.channel()
|
||
queue = await channel.get_queue("analyze.q", ensure=False)
|
||
deadline = 10.0
|
||
found = False
|
||
while deadline > 0:
|
||
start = time.monotonic()
|
||
try:
|
||
message = await queue.get(timeout=deadline)
|
||
except aio_pika.exceptions.QueueEmpty:
|
||
# basic_get is non-blocking; the routed message may not be
|
||
# visible yet. Treat as transient and retry until the deadline.
|
||
await asyncio.sleep(0.25)
|
||
deadline -= time.monotonic() - start
|
||
continue
|
||
await message.ack()
|
||
body = json.loads(message.body.decode("utf-8"))
|
||
if body["document_id"] == str(document_id):
|
||
assert body["correlation_id"] == str(correlation_id)
|
||
assert body["extracted_s3_key"] == extracted_key_path
|
||
assert body["char_count"] > 0
|
||
assert body["ocr_used"] is False
|
||
found = True
|
||
break
|
||
deadline -= time.monotonic() - start
|
||
assert found, "expected DocumentExtracted message not found in analyze.q"
|
||
finally:
|
||
await connection.close()
|
||
|
||
async with sess_factory() as session:
|
||
res = await session.execute(
|
||
text("SELECT status, stage, extracted_s3_key FROM documents WHERE id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
row = res.first()
|
||
assert row is not None
|
||
status, stage, ext_key = row
|
||
assert status == "analyzing"
|
||
assert stage == "queued_analyze"
|
||
assert ext_key == extracted_key_path
|
||
|
||
|
||
async def test_extract_worker_docx_uploads_text(
|
||
infra: dict[str, str],
|
||
docx_bytes: bytes,
|
||
) -> None:
|
||
from sqlalchemy import text
|
||
|
||
telegram_id = 444_555_666
|
||
sess_factory = session_factory(infra)
|
||
store = storage(infra)
|
||
async with sess_factory() 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()
|
||
|
||
document_id = uuid.uuid4()
|
||
correlation_id = uuid.uuid4()
|
||
s3_key = original_key(str(user_id), str(document_id), ".docx")
|
||
await store.put(
|
||
s3_key,
|
||
docx_bytes,
|
||
content_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||
)
|
||
|
||
async with sess_factory() as session:
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO documents (id, user_id, s3_key, filename, mime, bytes, status) "
|
||
"VALUES (:id, :uid, :s3, :fn, :mime, :bytes, 'queued')"
|
||
),
|
||
{
|
||
"id": document_id,
|
||
"uid": user_id,
|
||
"s3": s3_key,
|
||
"fn": "contract.docx",
|
||
"mime": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||
"bytes": len(docx_bytes),
|
||
},
|
||
)
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO jobs (document_id, correlation_id, queue, status) "
|
||
"VALUES (:did, :cid, 'extract', 'pending')"
|
||
),
|
||
{"did": document_id, "cid": correlation_id},
|
||
)
|
||
await session.commit()
|
||
|
||
handler = ExtractHandler(session_factory=sess_factory)
|
||
await handler.handle(
|
||
DocumentUploaded(
|
||
correlation_id=correlation_id,
|
||
document_id=document_id,
|
||
user_id=user_id,
|
||
s3_key=s3_key,
|
||
filename="contract.docx",
|
||
mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||
)
|
||
)
|
||
|
||
extracted_key_path = f"users/{user_id}/docs/{document_id}.txt"
|
||
text_data = await store.get(extracted_key_path)
|
||
assert "Стороны" in text_data.decode("utf-8")
|
||
|
||
async with sess_factory() as session:
|
||
row = await session.execute(
|
||
text("SELECT status FROM documents WHERE id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
assert row.scalar_one() == "analyzing"
|
||
|
||
|
||
async def test_extract_worker_terminal_failure_refunds_and_dlqs(
|
||
infra: dict[str, str],
|
||
) -> None:
|
||
from sqlalchemy import text
|
||
|
||
telegram_id = 777_888_999
|
||
sess_factory = session_factory(infra)
|
||
store = storage(infra)
|
||
async with sess_factory() as session:
|
||
result = await session.execute(
|
||
text(
|
||
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, 3) "
|
||
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 3 "
|
||
"RETURNING id, credits_left"
|
||
),
|
||
{"t": telegram_id},
|
||
)
|
||
row = result.first()
|
||
assert row is not None
|
||
user_id, initial_credits = row
|
||
await session.commit()
|
||
|
||
document_id = uuid.uuid4()
|
||
correlation_id = uuid.uuid4()
|
||
s3_key = original_key(str(user_id), str(document_id), ".pdf")
|
||
await store.put(s3_key, b"not a pdf", content_type="application/pdf")
|
||
|
||
async with sess_factory() as session:
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO documents (id, user_id, s3_key, filename, mime, bytes, status) "
|
||
"VALUES (:id, :uid, :s3, :fn, :mime, :bytes, 'queued')"
|
||
),
|
||
{
|
||
"id": document_id,
|
||
"uid": user_id,
|
||
"s3": s3_key,
|
||
"fn": "bad.pdf",
|
||
"mime": "application/pdf",
|
||
"bytes": 9,
|
||
},
|
||
)
|
||
await session.execute(
|
||
text(
|
||
"INSERT INTO jobs (document_id, correlation_id, queue, status, max_attempts) "
|
||
"VALUES (:did, :cid, 'extract', 'pending', 2)"
|
||
),
|
||
{"did": document_id, "cid": correlation_id},
|
||
)
|
||
await session.commit()
|
||
|
||
handler = ExtractHandler(session_factory=sess_factory)
|
||
from contract_check.core.analysis.extractor import ExtractionError
|
||
from contract_check.core.analysis.ocr import OCRError
|
||
|
||
# The handler may raise ExtractionError or OCRError depending on local
|
||
# tesseract data; either way the terminal-failure path refunds the credit.
|
||
with pytest.raises((ExtractionError, OCRError)):
|
||
await handler.handle(
|
||
DocumentUploaded(
|
||
correlation_id=correlation_id,
|
||
document_id=document_id,
|
||
user_id=user_id,
|
||
s3_key=s3_key,
|
||
filename="bad.pdf",
|
||
mime="application/pdf",
|
||
)
|
||
)
|
||
|
||
await handler.on_terminal_failure(
|
||
DocumentUploaded(
|
||
correlation_id=correlation_id,
|
||
document_id=document_id,
|
||
user_id=user_id,
|
||
s3_key=s3_key,
|
||
filename="bad.pdf",
|
||
mime="application/pdf",
|
||
),
|
||
"extraction_failed",
|
||
"test terminal failure",
|
||
)
|
||
|
||
async with sess_factory() as session:
|
||
res = await session.execute(
|
||
text("SELECT status, refunded FROM documents WHERE id = :d"),
|
||
{"d": document_id},
|
||
)
|
||
row = res.first()
|
||
assert row is not None
|
||
status, refunded = row
|
||
assert status == "failed"
|
||
assert refunded is True
|
||
|
||
credits = await session.execute(
|
||
text("SELECT credits_left FROM users WHERE id = :u"),
|
||
{"u": user_id},
|
||
)
|
||
assert credits.scalar_one() == initial_credits + 1
|