309 lines
13 KiB
Python
309 lines
13 KiB
Python
"""Общая обвязка unit-тестов воркера шлюза (`test_worker.py`, `test_worker_concurrency.py`); фикстура `make_env` —
|
||
в `tests/unit/conftest.py`.
|
||
|
||
Без Docker и сети: S3 — словарь в памяти (ошибки — классы `src.core.clients`), NATS — настоящий `NatsGateway` поверх
|
||
фейкового клиента nats-py (тело, `Nats-Msg-Id`, повторы и `PublishFailed` — как в проде), горячая папка — настоящий
|
||
`LocalHotFolder` на `tmp_path` под шпионом, ИС Сканирование — `process_once` эмулятора (задержка 0) в фоновой задаче.
|
||
Общий журнал событий (`events`) упорядочивает публикации, чтения S3 и запись во входящий каталог.
|
||
Каждый опубликованный конверт проверяется JSON Schema `recognition-result-message.schema.json`.
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import uuid
|
||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||
from contextlib import asynccontextmanager, suppress
|
||
from dataclasses import dataclass, field
|
||
from datetime import UTC, datetime, timedelta
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from jsonschema import Draft202012Validator, FormatChecker
|
||
from nats import errors as nats_errors
|
||
|
||
from src.contract import NATS_MSG_ID_HEADER, RESULT_SUBJECT
|
||
from src.core.clients import JSON_CONTENT_TYPE, OCTET_STREAM, S3ObjectNotFoundError, S3WriteError
|
||
from src.core.config import Settings
|
||
from src.core.nats import NatsGateway
|
||
from src.core.worker import Worker
|
||
from src.gateway.hotfolder import AnswerFile, LocalHotFolder
|
||
from src.scan_emulator import process_once
|
||
|
||
API_DIR = Path(__file__).resolve().parents[1]
|
||
CONTRACTS_DIR = API_DIR / "contracts"
|
||
BUCKET = "arm-ka"
|
||
PDF_BYTES = b"%PDF-1.4\n% scan-gateway test\n\x00\x01\x02binary"
|
||
|
||
MESSAGE_VALIDATOR = Draft202012Validator(
|
||
json.loads((CONTRACTS_DIR / "recognition-result-message.schema.json").read_text(encoding="utf-8")),
|
||
format_checker=FormatChecker(),
|
||
)
|
||
|
||
# Успешный ответ ИС в форме реальных `fixtures/ocr/aep_3.json`, `aep_6.json` (OQ-77): успех — `error.code == ""`,
|
||
# `error` — объект `{"code": "", "message": ""}`, а не `null`; `table_structure` — список таблиц с `label` и `rows`
|
||
SUCCESS_ANSWER: dict[str, Any] = {
|
||
"file_uid": "ok_doc",
|
||
"error": {"code": "", "message": ""},
|
||
"data_ocr": {
|
||
"label": {"form_name": "Бухгалтерский баланс", "confidence": 98.0},
|
||
"organization": {"form_name": "ООО Ромашка", "confidence": 97.5},
|
||
"inn": {"form_name": "7700000000", "confidence": 99.0},
|
||
"units": {"form_name": "тыс. руб.", "confidence": 95.0},
|
||
"date": {"form_name": "30.09.2025", "confidence": 96.0},
|
||
"table_structure": [
|
||
{
|
||
"label": {"form_name": "Бухгалтерский баланс", "confidence": 99.0},
|
||
"rows": [
|
||
{
|
||
"row_index": 0,
|
||
"cells": [
|
||
{"column_index": 0, "value": "БАЛАНС", "confidence": 97.0},
|
||
{"column_index": 1, "value": "1600", "confidence": 99.0},
|
||
{"column_index": 2, "value": "5678", "confidence": 96.0},
|
||
],
|
||
}
|
||
],
|
||
}
|
||
],
|
||
},
|
||
}
|
||
|
||
|
||
def merged_raw(job_id: str, *files: tuple[str, object, str | None]) -> dict[str, Any]:
|
||
"""Ожидаемый склеенный ответ `raw/{job_id}.json` (D-4): `(имя, response, parse_error)` по возрастанию имени."""
|
||
return {
|
||
"job_id": job_id,
|
||
"files": [{"name": name, "response": response, "parse_error": error} for name, response, error in files],
|
||
}
|
||
|
||
|
||
# --- фейки --------------------------------------------------------------------------------------------------------
|
||
class FakeStorage:
|
||
"""S3 в памяти с интерфейсом `S3Storage`: get_bytes / put_bytes / put_json и ошибками `src.core.clients`."""
|
||
|
||
def __init__(self, events: list[tuple[str, str]], bucket: str = BUCKET) -> None:
|
||
self.events = events
|
||
self.bucket = bucket
|
||
self.objects: dict[tuple[str, str], bytes] = {}
|
||
self.content_types: dict[tuple[str, str], str] = {}
|
||
self.get_calls: list[tuple[str, str]] = []
|
||
self.fail_put_prefixes: tuple[str, ...] = ()
|
||
|
||
async def get_bytes(self, key: str, *, bucket: str | None = None) -> bytes:
|
||
bucket_name = bucket or self.bucket
|
||
self.get_calls.append((bucket_name, key))
|
||
self.events.append(("s3_get", key))
|
||
try:
|
||
return self.objects[(bucket_name, key)]
|
||
except KeyError:
|
||
raise S3ObjectNotFoundError(
|
||
f"S3: объект {bucket_name}/{key} не прочитан (NoSuchKey)", bucket=bucket_name, key=key
|
||
) from None
|
||
|
||
async def put_bytes(self, key: str, data: bytes, content_type: str = OCTET_STREAM) -> None:
|
||
if self.fail_put_prefixes and key.startswith(self.fail_put_prefixes):
|
||
raise S3WriteError(
|
||
f"S3: объект {self.bucket}/{key} не записан (ConnectionError)", bucket=self.bucket, key=key
|
||
)
|
||
self.objects[(self.bucket, key)] = data
|
||
self.content_types[(self.bucket, key)] = content_type
|
||
self.events.append(("s3_put", key))
|
||
|
||
async def put_json(self, key: str, obj: Any) -> None:
|
||
await self.put_bytes(key, json.dumps(obj, ensure_ascii=False).encode("utf-8"), JSON_CONTENT_TYPE)
|
||
|
||
|
||
class FakeNatsClient:
|
||
"""Подмена `nats.aio.client.Client` для настоящего `NatsGateway`: записывает публикации в общий журнал."""
|
||
|
||
def __init__(self, events: list[tuple[str, str]]) -> None:
|
||
self.events = events
|
||
self.is_connected = False
|
||
self.is_closed = False
|
||
self.published: list[tuple[str, bytes, dict[str, str]]] = []
|
||
self.fail = False
|
||
self.block: asyncio.Event | None = None
|
||
|
||
async def connect(self, *args: Any, **kwargs: Any) -> None:
|
||
self.is_connected = True
|
||
|
||
async def publish(self, subject: str, payload: bytes = b"", headers: dict[str, str] | None = None) -> None:
|
||
if self.block is not None:
|
||
await self.block.wait()
|
||
if self.fail:
|
||
raise nats_errors.ConnectionClosedError
|
||
assert headers is not None
|
||
self.published.append((subject, payload, headers))
|
||
self.events.append(("publish", headers[NATS_MSG_ID_HEADER]))
|
||
|
||
async def flush(self, timeout: float = 10) -> None:
|
||
return None
|
||
|
||
async def drain(self) -> None:
|
||
self.is_connected = False
|
||
self.is_closed = True
|
||
|
||
async def close(self) -> None:
|
||
self.is_connected = False
|
||
self.is_closed = True
|
||
|
||
|
||
class SpyHotFolder:
|
||
"""`LocalHotFolder` под шпионом: журнал записей во входящий каталог, сценарии сбоев `poll`."""
|
||
|
||
def __init__(self, inner: LocalHotFolder, events: list[tuple[str, str]]) -> None:
|
||
self.inner = inner
|
||
self.events = events
|
||
self.puts: list[str] = []
|
||
self.polls = 0
|
||
self.poll_errors: list[BaseException] = []
|
||
self.poll_always: BaseException | None = None
|
||
self.poll_script: list[list[AnswerFile] | BaseException | None] = []
|
||
|
||
def put(self, name: str, data: bytes) -> None:
|
||
self.inner.put(name, data)
|
||
self.puts.append(name)
|
||
self.events.append(("put", name))
|
||
|
||
def poll(self, job_id: str) -> list[AnswerFile] | None:
|
||
self.polls += 1
|
||
if self.poll_always is not None:
|
||
raise self.poll_always
|
||
if self.poll_errors:
|
||
raise self.poll_errors.pop(0)
|
||
if self.poll_script:
|
||
step = self.poll_script.pop(0)
|
||
if isinstance(step, BaseException):
|
||
raise step
|
||
return step
|
||
return self.inner.poll(job_id)
|
||
|
||
def healthy(self) -> bool:
|
||
return self.inner.healthy()
|
||
|
||
|
||
# --- окружение теста ----------------------------------------------------------------------------------------------
|
||
@dataclass
|
||
class Env:
|
||
settings: Settings
|
||
root: Path
|
||
fixtures: Path
|
||
events: list[tuple[str, str]]
|
||
storage: FakeStorage
|
||
client: FakeNatsClient
|
||
gateway: NatsGateway
|
||
hotfolder: SpyHotFolder
|
||
worker: Worker
|
||
local: LocalHotFolder = field(init=False)
|
||
|
||
def __post_init__(self) -> None:
|
||
self.local = self.hotfolder.inner
|
||
|
||
# запрос -------------------------------------------------------------------------------------------------------
|
||
def request( # noqa: PLR0913 - параметры запроса для сценариев
|
||
self,
|
||
*,
|
||
original_name: str = "ok_doc.pdf",
|
||
deadline_s: float = 60,
|
||
attempt: int = 1,
|
||
bucket: str = BUCKET,
|
||
data: bytes | None = PDF_BYTES,
|
||
job_id: str | None = None,
|
||
) -> dict[str, Any]:
|
||
job = job_id or str(uuid.uuid4())
|
||
file_id = str(uuid.uuid4())
|
||
key = f"uploads/{file_id}/{job}/source.bin"
|
||
if data is not None:
|
||
self.storage.objects[(bucket, key)] = data
|
||
deadline = datetime.now(UTC) + timedelta(seconds=deadline_s)
|
||
return {
|
||
"schema_version": "1.0",
|
||
"job_id": job,
|
||
"file_id": file_id,
|
||
"pipeline": "image",
|
||
"file_type": "PDF",
|
||
"original_name": original_name,
|
||
"s3": {"bucket": bucket, "key": key},
|
||
"reply_deadline": deadline.isoformat().replace("+00:00", "Z"),
|
||
"attempt": attempt,
|
||
}
|
||
|
||
async def deliver(self, body: dict[str, Any] | bytes) -> None:
|
||
data = body if isinstance(body, bytes) else json.dumps(body, ensure_ascii=False).encode("utf-8")
|
||
await self.worker.on_request(FakeMsg(data))
|
||
|
||
async def run(self, body: dict[str, Any], timeout: float = 5) -> dict[str, Any]:
|
||
"""Доставить запрос, дождаться окончания всех заданий, вернуть единственный итоговый конверт задания."""
|
||
await self.deliver(body)
|
||
await asyncio.wait_for(self.worker.wait_idle(), timeout)
|
||
results = self.results(body["job_id"])
|
||
assert len(results) == 1, results
|
||
return results[0][1]
|
||
|
||
# публикации ---------------------------------------------------------------------------------------------------
|
||
def published(self) -> list[tuple[str, dict[str, Any]]]:
|
||
"""(Nats-Msg-Id, тело) всех публикаций; каждое тело проверено схемой конверта."""
|
||
out = []
|
||
for subject, payload, headers in self.client.published:
|
||
assert subject == RESULT_SUBJECT
|
||
body = json.loads(payload.decode("utf-8"))
|
||
MESSAGE_VALIDATOR.validate(body)
|
||
assert body["service"]["name"] == "scan-gateway"
|
||
assert body["schema_version"] == "1.0"
|
||
out.append((headers[NATS_MSG_ID_HEADER], body))
|
||
return out
|
||
|
||
def acks(self, job_id: str) -> list[str]:
|
||
return [mid for mid, body in self.published() if body["job_id"] == job_id and body["status"] == "processing"]
|
||
|
||
def results(self, job_id: str) -> list[tuple[str, dict[str, Any]]]:
|
||
return [
|
||
(mid, body) for mid, body in self.published() if body["job_id"] == job_id and body["status"] != "processing"
|
||
]
|
||
|
||
def inbox(self) -> list[str]:
|
||
if not self.local.input_dir.is_dir():
|
||
return []
|
||
return sorted(p.name for p in self.local.input_dir.iterdir() if not p.name.startswith("."))
|
||
|
||
def answer(self, job_id: str, content: bytes, name: str | None = None) -> None:
|
||
path = self.local.answer_dir(job_id) / (name or f"{job_id}.json")
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
path.write_bytes(content)
|
||
|
||
|
||
class FakeMsg:
|
||
def __init__(self, data: bytes) -> None:
|
||
self.data = data
|
||
self.subject = "recognition.image.request"
|
||
self.headers: dict[str, str] | None = None
|
||
|
||
|
||
EnvFactory = Callable[..., Awaitable[Env]]
|
||
|
||
|
||
@asynccontextmanager
|
||
async def emulator(env: Env, fixtures: Path | None = None) -> AsyncIterator[None]:
|
||
"""ИС Сканирование: `process_once` (задержка 0) каждые 5 мс, пока открыт контекст."""
|
||
|
||
async def loop() -> None:
|
||
while True:
|
||
process_once(env.root, fixtures or env.fixtures, 0)
|
||
await asyncio.sleep(0.005)
|
||
|
||
task = asyncio.create_task(loop())
|
||
try:
|
||
yield
|
||
finally:
|
||
task.cancel()
|
||
with suppress(asyncio.CancelledError):
|
||
await task
|
||
|
||
|
||
async def eventually(predicate: Callable[[], bool], timeout: float = 3.0) -> None:
|
||
deadline = asyncio.get_running_loop().time() + timeout
|
||
while not predicate():
|
||
if asyncio.get_running_loop().time() > deadline:
|
||
raise AssertionError("условие не выполнилось вовремя")
|
||
await asyncio.sleep(0.005)
|