209 lines
6.3 KiB
Python
209 lines
6.3 KiB
Python
"""B2B API integration tests: API-key auth, upload, report polling, management.
|
|
|
|
Run against the Docker Compose infrastructure (`docker compose up -d`).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import aio_pika
|
|
import httpx
|
|
import pytest
|
|
from tests.integration.conftest import user_token
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
@pytest.fixture
|
|
def pdf_bytes(tmp_path: Path) -> bytes:
|
|
import pymupdf
|
|
|
|
doc = pymupdf.open()
|
|
page = doc.new_page()
|
|
page.insert_text((72, 72), "Договор. Стороны обязуются.")
|
|
path = tmp_path / "contract.pdf"
|
|
doc.save(str(path))
|
|
doc.close()
|
|
return path.read_bytes()
|
|
|
|
|
|
@pytest.fixture
|
|
async def api_key_client(
|
|
client: httpx.AsyncClient,
|
|
db_session, # noqa: ANN001
|
|
infra: dict[str, str],
|
|
) -> tuple[httpx.AsyncClient, str, int]:
|
|
"""Create a user, authenticate, create an API key, and return (client, key, telegram_id)."""
|
|
from sqlalchemy import text
|
|
|
|
from contract_check.core.api_keys import hash_api_key
|
|
|
|
telegram_id = 999000999
|
|
await db_session.execute(
|
|
text(
|
|
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, 10) "
|
|
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 10"
|
|
),
|
|
{"t": telegram_id},
|
|
)
|
|
# Clean any leftover key from a previous interrupted run.
|
|
await db_session.execute(
|
|
text(
|
|
"DELETE FROM api_keys "
|
|
"WHERE user_id = (SELECT id FROM users WHERE telegram_id = :t) "
|
|
"AND name = :name"
|
|
),
|
|
{"t": telegram_id, "name": "integration-test-key"},
|
|
)
|
|
await db_session.commit()
|
|
|
|
token = await user_token(client, infra, telegram_id)
|
|
|
|
create_resp = await client.post(
|
|
"/api/v1/b2b/keys",
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
json={"name": "integration-test-key"},
|
|
)
|
|
assert create_resp.status_code == 201
|
|
body = create_resp.json()
|
|
raw_key = body["api_key"]
|
|
|
|
# Confirm hash is stored.
|
|
result = await db_session.execute(
|
|
text("SELECT id FROM api_keys WHERE key_hash = :h"),
|
|
{"h": hash_api_key(raw_key)},
|
|
)
|
|
assert result.first() is not None
|
|
|
|
return client, raw_key, telegram_id
|
|
|
|
|
|
async def test_b2b_analyze_upload_returns_202_and_enqueues(
|
|
api_key_client: tuple[httpx.AsyncClient, str, int],
|
|
pdf_bytes: bytes,
|
|
infra: dict[str, str],
|
|
) -> None:
|
|
client, api_key, _telegram_id = api_key_client
|
|
|
|
response = await client.post(
|
|
"/api/v1/analyze",
|
|
headers={"X-API-Key": api_key},
|
|
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"] == 9
|
|
|
|
# Verify a DocumentUploaded message landed on extract.q (drain leftovers).
|
|
connection = await aio_pika.connect_robust(infra["rabbitmq_url"])
|
|
try:
|
|
channel = await connection.channel()
|
|
queue = await channel.get_queue("extract.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()
|
|
msg_body = message.body.decode("utf-8")
|
|
if str(document_id) in msg_body:
|
|
assert str(correlation_id) in msg_body
|
|
found = True
|
|
break
|
|
deadline -= time.monotonic() - start
|
|
assert found, "expected DocumentUploaded message not found in extract.q"
|
|
finally:
|
|
await connection.close()
|
|
|
|
|
|
async def test_b2b_get_report_before_ready_returns_status(
|
|
api_key_client: tuple[httpx.AsyncClient, str, int],
|
|
pdf_bytes: bytes,
|
|
) -> None:
|
|
client, api_key, _telegram_id = api_key_client
|
|
|
|
upload = await client.post(
|
|
"/api/v1/analyze",
|
|
headers={"X-API-Key": api_key},
|
|
files={"file": ("contract.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
assert upload.status_code == 202
|
|
document_id = upload.json()["document_id"]
|
|
|
|
poll = await client.get(
|
|
f"/api/v1/b2b/reports/{document_id}",
|
|
headers={"X-API-Key": api_key},
|
|
)
|
|
assert poll.status_code == 200
|
|
body = poll.json()
|
|
assert body["document_id"] == document_id
|
|
assert body["status"] == "queued"
|
|
assert "stage" in body
|
|
|
|
|
|
async def test_b2b_missing_api_key_returns_401(client: httpx.AsyncClient) -> None:
|
|
response = await client.get("/api/v1/b2b/usage")
|
|
assert response.status_code == 401
|
|
|
|
|
|
async def test_b2b_invalid_api_key_returns_401(client: httpx.AsyncClient) -> None:
|
|
response = await client.get(
|
|
"/api/v1/b2b/usage",
|
|
headers={"X-API-Key": "clearly-invalid-key"},
|
|
)
|
|
assert response.status_code == 401
|
|
|
|
|
|
async def test_b2b_key_management_requires_user_jwt(
|
|
client: httpx.AsyncClient,
|
|
) -> None:
|
|
response = await client.post(
|
|
"/api/v1/b2b/keys",
|
|
json={"name": "no-auth-key"},
|
|
)
|
|
assert response.status_code == 401
|
|
|
|
|
|
async def test_b2b_revoke_key_blocks_usage(
|
|
api_key_client: tuple[httpx.AsyncClient, str, int],
|
|
infra: dict[str, str],
|
|
pdf_bytes: bytes,
|
|
) -> None:
|
|
client, api_key, telegram_id = api_key_client
|
|
|
|
token = await user_token(client, infra, telegram_id)
|
|
|
|
# List keys to find id.
|
|
list_resp = await client.get(
|
|
"/api/v1/b2b/keys",
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
assert list_resp.status_code == 200
|
|
key_id = list_resp.json()[0]["id"]
|
|
|
|
revoke = await client.post(
|
|
f"/api/v1/b2b/keys/{key_id}/revoke",
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
)
|
|
assert revoke.status_code == 200
|
|
|
|
# Revoked key cannot upload.
|
|
response = await client.post(
|
|
"/api/v1/analyze",
|
|
headers={"X-API-Key": api_key},
|
|
files={"file": ("contract.pdf", pdf_bytes, "application/pdf")},
|
|
)
|
|
assert response.status_code == 401
|