197 lines
7.1 KiB
Python
197 lines
7.1 KiB
Python
"""Upload pipeline integration test: POST /api/v1/documents → MinIO + RabbitMQ + credits.
|
||
|
||
Verifies the api:
|
||
- accepts a multipart PDF upload with a service token and telegram_id,
|
||
- reserves a credit,
|
||
- stores the blob in MinIO under the expected key,
|
||
- publishes a `DocumentUploaded` message to the `extract.q` queue,
|
||
- returns 202 with document_id/correlation_id/credits_left.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import TYPE_CHECKING
|
||
|
||
import aio_pika
|
||
import httpx
|
||
import pytest
|
||
from asgi_lifespan import LifespanManager
|
||
|
||
from contract_check.api.app import create_app
|
||
from tests.integration.conftest import user_token
|
||
|
||
if TYPE_CHECKING:
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
pytestmark = pytest.mark.integration
|
||
|
||
|
||
@pytest.fixture
|
||
def pdf_bytes(tmp_path: Path) -> bytes:
|
||
# The real test requires pymupdf to create a minimal PDF.
|
||
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()
|
||
|
||
|
||
async def test_upload_document_reserves_credit_and_enqueues(
|
||
client: httpx.AsyncClient,
|
||
infra: dict[str, str],
|
||
pdf_bytes: bytes,
|
||
db_session: AsyncSession,
|
||
) -> None:
|
||
from sqlalchemy import text
|
||
|
||
telegram_id = 123456789
|
||
await db_session.execute(
|
||
text(
|
||
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, 5) "
|
||
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 5"
|
||
),
|
||
{"t": telegram_id},
|
||
)
|
||
await db_session.commit()
|
||
|
||
token = await user_token(client, infra, telegram_id)
|
||
|
||
response = await client.post( # type: ignore[misc]
|
||
"/api/v1/documents",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
files={"file": ("contract.pdf", pdf_bytes, "application/pdf")},
|
||
)
|
||
assert response.status_code == 202
|
||
body = response.json()
|
||
document_id = uuid.UUID(body["document_id"])
|
||
correlation_id = uuid.UUID(body["correlation_id"])
|
||
assert body["credits_left"] == 4 # 5 - 1
|
||
|
||
# Verify a DocumentUploaded message landed on extract.q.
|
||
# When worker-extract is running it consumes quickly; in that case we
|
||
# verify the pipeline completes by checking analyze.q instead.
|
||
import time
|
||
|
||
connection = await aio_pika.connect_robust(infra["rabbitmq_url"])
|
||
try:
|
||
channel = await connection.channel()
|
||
deadline = 10.0
|
||
found = False
|
||
try:
|
||
queue = await channel.get_queue("extract.q", ensure=False)
|
||
while deadline > 0:
|
||
start = time.monotonic()
|
||
message = await queue.get(timeout=deadline)
|
||
await message.ack()
|
||
msg_body = json.loads(message.body.decode("utf-8"))
|
||
if msg_body["document_id"] == str(document_id):
|
||
assert msg_body["correlation_id"] == str(correlation_id)
|
||
assert msg_body["user_id"] # any valid UUID
|
||
assert msg_body["s3_key"].endswith(".pdf")
|
||
assert msg_body["mime"] == "application/pdf"
|
||
found = True
|
||
break
|
||
deadline -= time.monotonic() - start
|
||
except aio_pika.exceptions.QueueEmpty:
|
||
# Worker already consumed it; assert end-to-end completion.
|
||
analyze_queue = await channel.get_queue("analyze.q", ensure=False)
|
||
# Use a consumer iterator so we don't depend on quorum-queue
|
||
# basic_get behaviour returning stale/empty results under concurrency.
|
||
async with analyze_queue.iterator() as queue_iter:
|
||
async for message in queue_iter:
|
||
async with message.process():
|
||
msg_body = json.loads(message.body.decode("utf-8"))
|
||
if msg_body["document_id"] == str(document_id):
|
||
assert msg_body["correlation_id"] == str(correlation_id)
|
||
assert msg_body["extracted_s3_key"].endswith(".txt")
|
||
assert msg_body["char_count"] > 0
|
||
found = True
|
||
break
|
||
if found:
|
||
break
|
||
if time.monotonic() - start > deadline:
|
||
break
|
||
assert found, "expected pipeline message not found on extract.q or analyze.q"
|
||
finally:
|
||
await connection.close()
|
||
|
||
|
||
async def test_upload_rejected_before_storage_on_billing_hold(
|
||
client: httpx.AsyncClient,
|
||
infra: dict[str, str],
|
||
pdf_bytes: bytes,
|
||
db_session: AsyncSession,
|
||
) -> None:
|
||
from sqlalchemy import text
|
||
|
||
telegram_id = 123456790
|
||
await db_session.execute(
|
||
text(
|
||
"INSERT INTO users (telegram_id, credits_left, billing_hold) VALUES (:t, 5, TRUE) "
|
||
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 5, billing_hold = TRUE"
|
||
),
|
||
{"t": telegram_id},
|
||
)
|
||
await db_session.commit()
|
||
|
||
token = await user_token(client, infra, telegram_id)
|
||
|
||
response = await client.post(
|
||
"/api/v1/documents",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
files={"file": ("contract.pdf", pdf_bytes, "application/pdf")},
|
||
)
|
||
assert response.status_code == 402
|
||
assert "billing hold" in response.text.lower()
|
||
|
||
|
||
async def test_upload_rejects_oversized_file(
|
||
infra: dict[str, str],
|
||
db_session: AsyncSession,
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
from sqlalchemy import text
|
||
|
||
telegram_id = 123456791
|
||
await db_session.execute(
|
||
text(
|
||
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, 5) "
|
||
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 5"
|
||
),
|
||
{"t": telegram_id},
|
||
)
|
||
await db_session.commit()
|
||
|
||
monkeypatch.setenv("MAX_UPLOAD_BYTES", "10")
|
||
from contract_check.core.config import get_settings
|
||
|
||
get_settings.cache_clear()
|
||
app = create_app()
|
||
async with httpx.AsyncClient(
|
||
transport=httpx.ASGITransport(app=app), base_url="http://test"
|
||
) as client:
|
||
async with LifespanManager(app):
|
||
token = await user_token(client, infra, telegram_id)
|
||
response = await client.post(
|
||
"/api/v1/documents",
|
||
headers={"Authorization": f"Bearer {token}"},
|
||
files={"file": ("big.pdf", b"x" * 100, "application/pdf")},
|
||
)
|
||
assert response.status_code == 413
|