157 lines
6.3 KiB
Python
157 lines
6.3 KiB
Python
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
from uuid import uuid4
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
from src.schemas.pdfworker import PDFWorkerCommand, PDFWorkerResult
|
|
from src.core.pdfworker_tasks import create_pdfworker_task
|
|
from src.core.pdfworker_nats import handle_result_message, publish_pdfworker_task, pdfworker_connection, nats_auth
|
|
|
|
|
|
def command():
|
|
return create_pdfworker_task('uploads/2026/report.pdf', uuid4(), task_id=uuid4(),
|
|
expires_at=datetime.now(timezone.utc) + timedelta(minutes=10))
|
|
|
|
|
|
def event(status='ready'):
|
|
task = command()
|
|
return PDFWorkerResult(task_id=task.task_id, file_id=task.file_id, status=status,
|
|
result_s3_key='pdfworker/results/result.json' if status == 'ready' else None,
|
|
error_code=None if status == 'ready' else 'storage_unavailable',
|
|
finished_at=datetime.now(timezone.utc))
|
|
|
|
|
|
def message(data):
|
|
return SimpleNamespace(data=data)
|
|
|
|
|
|
def test_full_key_and_persisted_identity():
|
|
a = command()
|
|
b = create_pdfworker_task(a.source_s3_key, a.file_id, task_id=a.task_id, expires_at=a.expires_at)
|
|
assert a == b
|
|
assert a.source_s3_key == 'uploads/2026/report.pdf'
|
|
assert 'version' not in a.model_dump()
|
|
assert 'legal_entity_id' not in a.model_dump()
|
|
assert PDFWorkerCommand.model_validate_json(a.model_dump_json()) == a
|
|
|
|
|
|
@pytest.mark.parametrize('extension', ['docx', 'zip', 'jpeg', 'csv', 'c++', 'данные'])
|
|
def test_arbitrary_safe_extension(extension):
|
|
task = command()
|
|
data = {**task.model_dump(), 'source_s3_key': f'uploads/report.{extension}'}
|
|
assert PDFWorkerCommand.model_validate(data).source_s3_key == data['source_s3_key']
|
|
|
|
|
|
@pytest.mark.parametrize('field,value', [('version', 2), ('task_id', 'bad'), ('file_id', 'bad'),
|
|
('expires_at', 'test'), ('expires_at', '2026-10-01T12:00:00'),
|
|
('source_s3_key', 's3://bucket/doc.pdf'), ('source_s3_key', ''),
|
|
('source_s3_key', 'uploads/report.bad:name'), ('source_files', ['x.json'])])
|
|
def test_invalid_command(field, value):
|
|
with pytest.raises(ValidationError):
|
|
PDFWorkerCommand.model_validate({**command().model_dump(), field: value})
|
|
|
|
|
|
@pytest.mark.parametrize('updates', [dict(version=2), dict(status='success'), dict(error_code='unknown'),
|
|
dict(result_s3_key=None), dict(error_code='storage_unavailable'),
|
|
dict(status='timeout', result_s3_key=None, error_code='source_not_found'),
|
|
dict(status='error', result_s3_key=None, error_code='deadline_exceeded')])
|
|
def test_invalid_result(updates):
|
|
with pytest.raises(ValidationError):
|
|
PDFWorkerResult.model_validate({**event().model_dump(), **updates})
|
|
|
|
|
|
def test_expired_command_still_valid_for_timeout_response():
|
|
data = command().model_dump()
|
|
data['expires_at'] = datetime.now(timezone.utc) - timedelta(days=1)
|
|
assert PDFWorkerCommand.model_validate(data)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_publish_command_and_flush():
|
|
nc = SimpleNamespace(publish=AsyncMock(), flush=AsyncMock())
|
|
task = command()
|
|
await publish_pdfworker_task(nc, task)
|
|
assert nc.publish.call_args.args[0] == 'files.tabular_uploaded'
|
|
assert PDFWorkerCommand.model_validate_json(nc.publish.call_args.args[1]) == task
|
|
nc.flush.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('status', ['ready', 'error'])
|
|
async def test_result_handler(status):
|
|
calls = []
|
|
msg = message(event(status).model_dump_json().encode())
|
|
async def handle(result):
|
|
assert result.status == status
|
|
calls.append('commit')
|
|
await handle_result_message(msg, handle)
|
|
assert calls == ['commit']
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invalid_message_is_ignored():
|
|
msg, handler = message(b'{}'), AsyncMock()
|
|
await handle_result_message(msg, handler)
|
|
handler.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handler_failure_is_logged():
|
|
msg = message(event().model_dump_json().encode())
|
|
await handle_result_message(msg, AsyncMock(side_effect=RuntimeError('db down')))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize('fail_subscribe', [False, True])
|
|
async def test_lifecycle_cleanup(monkeypatch, fail_subscribe):
|
|
from src.core import pdfworker_nats
|
|
sub = SimpleNamespace(drain=AsyncMock())
|
|
subscribe = AsyncMock(return_value=sub)
|
|
if fail_subscribe:
|
|
subscribe.side_effect = RuntimeError('subscription denied')
|
|
nc = SimpleNamespace(is_closed=False, is_connected=True, connect=AsyncMock(),
|
|
subscribe=subscribe, flush=AsyncMock(), drain=AsyncMock(), close=AsyncMock())
|
|
monkeypatch.setattr(pdfworker_nats, 'NATS', lambda: nc)
|
|
if fail_subscribe:
|
|
with pytest.raises(RuntimeError):
|
|
async with pdfworker_connection('nats://localhost', AsyncMock()):
|
|
pytest.fail('should not yield')
|
|
else:
|
|
async with pdfworker_connection('nats://localhost', AsyncMock()) as connected:
|
|
assert connected is nc
|
|
pass
|
|
subscribe.assert_awaited_once()
|
|
sub.drain.assert_awaited_once()
|
|
nc.drain.assert_awaited_once()
|
|
nc.close.assert_awaited_once()
|
|
|
|
|
|
def test_optional_auth():
|
|
assert nats_auth() == {}
|
|
with pytest.raises(ValueError):
|
|
nats_auth('jwt', '')
|
|
|
|
|
|
def test_models_match_pdfworker_checkout():
|
|
counterpart = Path(__file__).resolve().parents[4] / 'DKB_ARMKA_FILE_FLOW/api/src/schemas/messages.py'
|
|
if not counterpart.exists():
|
|
pytest.skip('PDFWorker checkout not available')
|
|
from runpy import run_path
|
|
models = run_path(str(counterpart))
|
|
assert PDFWorkerCommand.model_json_schema() == models['PDFWorkerCommand'].model_json_schema()
|
|
assert PDFWorkerResult.model_json_schema() == models['PDFWorkerResult'].model_json_schema()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_failed_connect_closes_allocated_client(monkeypatch):
|
|
from src.core import pdfworker_nats
|
|
nc = SimpleNamespace(is_closed=False, is_connected=False,
|
|
connect=AsyncMock(side_effect=RuntimeError('offline')), close=AsyncMock())
|
|
monkeypatch.setattr(pdfworker_nats, 'NATS', lambda: nc)
|
|
with pytest.raises(RuntimeError, match='offline'):
|
|
async with pdfworker_connection('nats://localhost', AsyncMock()):
|
|
pytest.fail('should not yield')
|
|
nc.close.assert_awaited_once()
|