feat: воркер распознавания src/core/worker.py — ack до чтения S3, дедуп job_id, таблица ошибок D-5
on_request: валидация (pydantic + MAJOR = 1) → повторный job_id получает только ack → ack processing → очередь; рабочие корутины: S3 get → Parser через asyncio.to_thread → файл результата → конверт с metrics.duration_ms → публикация с Nats-Msg-Id result:<job_id>:<attempt>. SCHEMA_MISMATCH без ack, FILE_READ_ERROR, UNKNOWN_ERROR, NatsPublishError — error-лог. Docstring протокола Parser: ожидаемые сбои — исход error (T7.1). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Andrey Epifancev
parent
ff3ec5684a
commit
e1231dcdb4
445
api/src/core/worker.py
Normal file
445
api/src/core/worker.py
Normal file
@ -0,0 +1,445 @@
|
||||
"""Воркер распознавания: запрос `recognition.tabular.request` → ack → S3 → парсер → S3 → `recognition.result`.
|
||||
|
||||
Design D-2, D-4, D-5, D-7; spec tabular-recognition-service «Подтверждение запроса до обработки», «Результат по
|
||||
контракту», «Ошибки транспорта превращаются в результат», «Парсер — заменяемый интерфейс».
|
||||
|
||||
- `on_request(msg)` — обработчик подписки NATS (queue group `tabular-parser`): разбор тела и валидация по контракту
|
||||
(pydantic `RecognitionRequest`, MAJOR `schema_version` = 1). Невалидный запрос — конверт `error SCHEMA_MISMATCH`
|
||||
без ack (если `job_id`, `file_id` и `attempt` читаемы; иначе собрать валидный конверт нельзя — error-лог и сброс).
|
||||
Валидный: `job_id` уже в работе или обработан этим процессом — только повторный ack (`ack:<job_id>:<attempt>` этого
|
||||
запроса), второй обработки нет; иначе — `job_id` в множество «в работе», ack `processing` **до** чтения S3, задание
|
||||
во внутреннюю `asyncio.Queue`.
|
||||
- Рабочие корутины (`concurrency`): S3 get → `Parser.parse` через `asyncio.to_thread` (R-13) → для `done | needs_check`
|
||||
файл результата `results/{file_id}/{job_id}/result.json` → конверт `done | needs_check | error` с
|
||||
`metrics.duration_ms` (всегда, ≥ 0) → публикация с `Nats-Msg-Id = result:<job_id>:<attempt>`.
|
||||
- Таблица ошибок D-5: S3 get упал → `FILE_READ_ERROR`; исключение парсера → `UNKNOWN_ERROR` (тип и текст исключения
|
||||
одной строкой, без трассировки); запись файла результата упала → `UNKNOWN_ERROR`; публикация не удалась после
|
||||
повторов (`NatsPublishError`) → error-лог, задание дальше не идёт (ядро закроет его по `reply_deadline`).
|
||||
Никакое исключение не роняет процесс и рабочие корутины.
|
||||
- Дедупликация: `job_id` в работе (очередь + обрабатываемые) — множество; завершённые (любой итог, в том числе ошибка и
|
||||
неудачная публикация результата) — LRU размера `done_cache_size`.
|
||||
- Остановка: новые запросы не принимаются (без ack — ядро переотправит другой инстанции), очередь и текущие задания
|
||||
дорабатываются не дольше `timeout`, остальное отменяется с warning-логом.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol
|
||||
|
||||
from loguru import logger
|
||||
from pydantic import ValidationError
|
||||
|
||||
from src.contract import (
|
||||
CONTRACT_VERSION,
|
||||
NATS_MSG_ID_HEADER,
|
||||
RESULT_SUBJECT,
|
||||
ErrorCode,
|
||||
MessageError,
|
||||
MessageStatus,
|
||||
Metrics,
|
||||
RecognitionRequest,
|
||||
RecognitionResult,
|
||||
RecognitionResultMessage,
|
||||
ResultStatus,
|
||||
S3Ref,
|
||||
ServiceInfo,
|
||||
ack_msg_id,
|
||||
result_msg_id,
|
||||
)
|
||||
from src.core.clients import Storage
|
||||
from src.core.nats import NatsPublishError, Publisher
|
||||
from src.parser import ParseOutcome, Parser
|
||||
|
||||
SUPPORTED_MAJOR = "1"
|
||||
STOP_TIMEOUT = 10.0 # секунд на доработку очереди при остановке (lifespan)
|
||||
ERROR_MESSAGE_LIMIT = 1000 # символов в `error.message` конверта
|
||||
|
||||
_VERSION_RE = re.compile(r"(\d+)\.\d+", re.ASCII)
|
||||
_UUID_RE = re.compile(r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}", re.ASCII)
|
||||
|
||||
Clock = Callable[[], float]
|
||||
|
||||
|
||||
class InboundMessage(Protocol):
|
||||
"""Входящее сообщение NATS (подмножество `nats.aio.msg.Msg`, нужное воркеру)."""
|
||||
|
||||
@property
|
||||
def data(self) -> bytes: ...
|
||||
|
||||
@property
|
||||
def headers(self) -> Mapping[str, str] | None: ...
|
||||
|
||||
|
||||
def result_key(file_id: object, job_id: object) -> str:
|
||||
"""Ключ файла результата в бакете запроса (design D-4)."""
|
||||
return f"results/{file_id}/{job_id}/result.json"
|
||||
|
||||
|
||||
class InvalidRequestError(Exception):
|
||||
"""Запрос не прошёл контракт: `SCHEMA_MISMATCH`."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Job:
|
||||
request: RecognitionRequest
|
||||
|
||||
|
||||
def _one_line(text: str) -> str:
|
||||
"""Сообщение ошибки для конверта: одна строка, без трассировки, ограниченной длины."""
|
||||
flat = " ".join(text.split())
|
||||
return flat if len(flat) <= ERROR_MESSAGE_LIMIT else flat[: ERROR_MESSAGE_LIMIT - 1] + "…"
|
||||
|
||||
|
||||
def _describe(exc: BaseException) -> str:
|
||||
return _one_line(f"{type(exc).__name__}: {exc}")
|
||||
|
||||
|
||||
def _validation_summary(exc: ValidationError) -> str:
|
||||
parts = []
|
||||
for error in exc.errors(include_url=False, include_input=False):
|
||||
location = ".".join(str(item) for item in error["loc"]) or "<тело>"
|
||||
parts.append(f"{location}: {error['msg']}")
|
||||
return "; ".join(parts)
|
||||
|
||||
|
||||
class RecognitionWorker:
|
||||
"""Обработчик запросов распознавания; зависимости — протоколы (`Storage`, `Publisher`, `Parser`, design D-7)."""
|
||||
|
||||
def __init__( # noqa: PLR0913 - все параметры, кроме зависимостей, именованные
|
||||
self,
|
||||
*,
|
||||
storage: Storage,
|
||||
publisher: Publisher,
|
||||
parser: Parser,
|
||||
service_name: str,
|
||||
service_version: str,
|
||||
concurrency: int = 1,
|
||||
done_cache_size: int = 1000,
|
||||
clock: Clock = time.monotonic,
|
||||
) -> None:
|
||||
if concurrency < 1 or done_cache_size < 1:
|
||||
raise ValueError("concurrency и done_cache_size должны быть ≥ 1")
|
||||
self._storage = storage
|
||||
self._publisher = publisher
|
||||
self._parser = parser
|
||||
self._service = ServiceInfo.model_validate({"name": service_name, "version": service_version})
|
||||
self._concurrency = concurrency
|
||||
self._done_cache_size = done_cache_size
|
||||
self._clock = clock
|
||||
self._queue: asyncio.Queue[_Job] = asyncio.Queue()
|
||||
self._in_progress: set[str] = set() # job_id в очереди или в обработке
|
||||
self._done: OrderedDict[str, None] = OrderedDict() # LRU завершённых job_id
|
||||
self._tasks: list[asyncio.Task[None]] = []
|
||||
self._accepting = True
|
||||
|
||||
# --- жизненный цикл ---------------------------------------------------------------------------------------------
|
||||
def start(self) -> None:
|
||||
"""Запускает рабочие корутины (`concurrency` штук)."""
|
||||
if self._tasks:
|
||||
return
|
||||
self._accepting = True
|
||||
self._tasks = [
|
||||
asyncio.create_task(self._run(), name=f"recognition-worker-{index}") for index in range(self._concurrency)
|
||||
]
|
||||
logger.info("Воркер распознавания запущен. concurrency={concurrency}", concurrency=self._concurrency)
|
||||
|
||||
async def join(self) -> None:
|
||||
"""Ждёт, пока очередь опустеет и все взятые задания завершатся."""
|
||||
await self._queue.join()
|
||||
|
||||
async def stop(self, timeout: float = STOP_TIMEOUT) -> None:
|
||||
"""Перестаёт принимать запросы, дорабатывает очередь не дольше `timeout` секунд, затем отменяет остальное."""
|
||||
self._accepting = False
|
||||
if not self._tasks:
|
||||
return
|
||||
try:
|
||||
await asyncio.wait_for(self._queue.join(), timeout)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"Остановка воркера: задания не завершились за отведённое время и отменены (ack опубликован, ядро "
|
||||
"закроет их по reply_deadline). error_category=shutdown, in_flight={in_flight}, pending={pending}",
|
||||
in_flight=len(self._in_progress) - self._queue.qsize(),
|
||||
pending=self._queue.qsize(),
|
||||
)
|
||||
tasks, self._tasks = self._tasks, []
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
logger.info("Воркер распознавания остановлен")
|
||||
|
||||
# --- приём запроса ----------------------------------------------------------------------------------------------
|
||||
async def on_request(self, msg: InboundMessage) -> None:
|
||||
"""Обработчик подписки на `recognition.tabular.request`; исключения наружу не выходят."""
|
||||
if not self._accepting:
|
||||
logger.warning("Воркер останавливается: запрос не принят (без ack, ядро переотправит)")
|
||||
return
|
||||
started = self._clock()
|
||||
try:
|
||||
request = self._parse_request(msg.data)
|
||||
except InvalidRequestError as exc:
|
||||
await self._reject(msg, str(exc), started)
|
||||
return
|
||||
|
||||
job_id = str(request.job_id)
|
||||
duplicate = job_id in self._in_progress or job_id in self._done
|
||||
if duplicate:
|
||||
if job_id in self._done:
|
||||
self._done.move_to_end(job_id)
|
||||
logger.info(
|
||||
"Повторный запрос: только ack, без повторной обработки. job_id={job_id}, attempt={attempt}",
|
||||
job_id=job_id,
|
||||
attempt=request.attempt,
|
||||
)
|
||||
else:
|
||||
self._in_progress.add(job_id)
|
||||
await self._publish_ack(request)
|
||||
if not duplicate:
|
||||
self._queue.put_nowait(_Job(request))
|
||||
|
||||
def _parse_request(self, data: bytes) -> RecognitionRequest:
|
||||
try:
|
||||
raw = json.loads(data)
|
||||
except (ValueError, UnicodeDecodeError) as exc:
|
||||
raise InvalidRequestError(f"тело запроса — не JSON ({type(exc).__name__})") from None
|
||||
if isinstance(raw, dict):
|
||||
version = raw.get("schema_version")
|
||||
match = _VERSION_RE.fullmatch(version) if isinstance(version, str) else None
|
||||
if match is not None and match.group(1) != SUPPORTED_MAJOR:
|
||||
raise InvalidRequestError(
|
||||
f"неподдерживаемая версия запроса schema_version={version!r}: сервис поддерживает "
|
||||
f"{SUPPORTED_MAJOR}.x"
|
||||
)
|
||||
try:
|
||||
return RecognitionRequest.model_validate(raw)
|
||||
except ValidationError as exc:
|
||||
raise InvalidRequestError(
|
||||
f"запрос не соответствует recognition-request.schema.json: {_validation_summary(exc)}"
|
||||
) from None
|
||||
|
||||
async def _reject(self, msg: InboundMessage, reason: str, started: float) -> None:
|
||||
"""`SCHEMA_MISMATCH` без ack. Без читаемых `job_id`/`file_id`/`attempt` конверт не собрать — лог и сброс."""
|
||||
raw = self._raw_dict(msg.data)
|
||||
job_id = self._readable_uuid(raw, "job_id")
|
||||
file_id = self._readable_uuid(raw, "file_id")
|
||||
attempt = self._readable_attempt(raw, msg.headers, job_id) if job_id is not None else None
|
||||
if job_id is None or file_id is None or attempt is None:
|
||||
logger.error(
|
||||
"Запрос отброшен: не проходит контракт, а job_id/file_id/attempt не читаемы — результат не "
|
||||
"публикуется. error_category=schema_mismatch, reason={reason}, size={size}",
|
||||
reason=reason,
|
||||
size=len(msg.data),
|
||||
)
|
||||
return
|
||||
logger.error(
|
||||
"Запрос не проходит контракт: публикуется SCHEMA_MISMATCH без ack. error_category=schema_mismatch, "
|
||||
"job_id={job_id}, reason={reason}",
|
||||
job_id=job_id,
|
||||
reason=reason,
|
||||
)
|
||||
envelope = self._envelope(
|
||||
job_id=job_id,
|
||||
file_id=file_id,
|
||||
status=MessageStatus.ERROR,
|
||||
error=MessageError(code=ErrorCode.SCHEMA_MISMATCH, message=_one_line(reason)),
|
||||
metrics=self._metrics(None, started),
|
||||
)
|
||||
await self._publish(envelope, result_msg_id(job_id, attempt), job_id)
|
||||
|
||||
@staticmethod
|
||||
def _raw_dict(data: bytes) -> dict[str, Any]:
|
||||
try:
|
||||
raw = json.loads(data)
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return {}
|
||||
return raw if isinstance(raw, dict) else {}
|
||||
|
||||
@staticmethod
|
||||
def _readable_uuid(raw: Mapping[str, Any], name: str) -> str | None:
|
||||
value = raw.get(name)
|
||||
return value if isinstance(value, str) and _UUID_RE.fullmatch(value) else None
|
||||
|
||||
@staticmethod
|
||||
def _readable_attempt(raw: Mapping[str, Any], headers: Mapping[str, str] | None, job_id: str) -> int | None:
|
||||
"""`attempt` из тела (целое ≥ 1), иначе из заголовка запроса `Nats-Msg-Id = <job_id>:<attempt>`."""
|
||||
attempt = raw.get("attempt")
|
||||
if isinstance(attempt, int) and not isinstance(attempt, bool) and attempt >= 1:
|
||||
return attempt
|
||||
header = (headers or {}).get(NATS_MSG_ID_HEADER, "")
|
||||
prefix, _, tail = header.rpartition(":")
|
||||
if prefix.lower() == job_id.lower() and tail.isdigit() and int(tail) >= 1:
|
||||
return int(tail)
|
||||
return None
|
||||
|
||||
async def _publish_ack(self, request: RecognitionRequest) -> None:
|
||||
envelope = self._envelope(job_id=request.job_id, file_id=request.file_id, status=MessageStatus.PROCESSING)
|
||||
await self._publish(envelope, ack_msg_id(request.job_id, request.attempt), str(request.job_id))
|
||||
|
||||
# --- обработка задания ------------------------------------------------------------------------------------------
|
||||
async def _run(self) -> None:
|
||||
while True:
|
||||
job = await self._queue.get()
|
||||
job_id = str(job.request.job_id)
|
||||
try:
|
||||
await self._process(job.request)
|
||||
except Exception:
|
||||
# Страховка: _process сам превращает ошибки в конверты; сюда попадает только непредвиденное
|
||||
logger.exception("Сбой обработки задания. error_category=worker_error, job_id={job_id}", job_id=job_id)
|
||||
finally:
|
||||
self._finish(job_id)
|
||||
self._queue.task_done()
|
||||
|
||||
def _finish(self, job_id: str) -> None:
|
||||
self._in_progress.discard(job_id)
|
||||
self._done[job_id] = None
|
||||
self._done.move_to_end(job_id)
|
||||
while len(self._done) > self._done_cache_size:
|
||||
self._done.popitem(last=False)
|
||||
|
||||
async def _process(self, request: RecognitionRequest) -> None:
|
||||
started = self._clock()
|
||||
job_id = str(request.job_id)
|
||||
envelope = await self._build_result(request, started)
|
||||
logger.info(
|
||||
"Задание обработано. job_id={job_id}, status={status}, error_code={error_code}",
|
||||
job_id=job_id,
|
||||
status=envelope.status.value,
|
||||
error_code=envelope.error.code.value if envelope.error else None,
|
||||
)
|
||||
await self._publish(envelope, result_msg_id(request.job_id, request.attempt), job_id)
|
||||
|
||||
async def _build_result(self, request: RecognitionRequest, started: float) -> RecognitionResultMessage:
|
||||
"""Конвейер S3 → парсер → S3; любой сбой — конверт `error` по таблице D-5."""
|
||||
try:
|
||||
data = await self._storage.get_bytes(request.s3.bucket, request.s3.key)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Не удалось прочитать исходный файл. error_category=s3_error, job_id={job_id}, key={key}, "
|
||||
"error={error}",
|
||||
job_id=request.job_id,
|
||||
key=request.s3.key,
|
||||
error=_describe(exc),
|
||||
)
|
||||
message = f"не удалось прочитать s3://{request.s3.bucket}/{request.s3.key}: {_describe(exc)}"
|
||||
return self._error(request, ErrorCode.FILE_READ_ERROR, message, None, started)
|
||||
|
||||
try:
|
||||
outcome: ParseOutcome = await asyncio.to_thread(self._parser.parse, data, request.file_type)
|
||||
except Exception as exc:
|
||||
logger.opt(exception=exc).error(
|
||||
"Исключение парсера. error_category=parser_error, job_id={job_id}", job_id=request.job_id
|
||||
)
|
||||
return self._error(request, ErrorCode.UNKNOWN_ERROR, f"ошибка разбора: {_describe(exc)}", None, started)
|
||||
|
||||
if outcome.status is ResultStatus.ERROR:
|
||||
error = outcome.error or MessageError(
|
||||
code=ErrorCode.UNKNOWN_ERROR, message="парсер вернул ошибку без описания"
|
||||
)
|
||||
return self._envelope(
|
||||
job_id=request.job_id,
|
||||
file_id=request.file_id,
|
||||
status=MessageStatus.ERROR,
|
||||
error=error,
|
||||
metrics=self._metrics(outcome.metrics, started),
|
||||
)
|
||||
|
||||
key = result_key(request.file_id, request.job_id)
|
||||
try:
|
||||
metrics = self._metrics(outcome.metrics, started)
|
||||
result = RecognitionResult.model_validate(
|
||||
{
|
||||
"schema_version": CONTRACT_VERSION,
|
||||
"job_id": request.job_id,
|
||||
"file_id": request.file_id,
|
||||
"service": self._service,
|
||||
"status": ResultStatus(outcome.status),
|
||||
"metrics": metrics,
|
||||
"documents": list(outcome.documents),
|
||||
}
|
||||
)
|
||||
await self._storage.put_json(request.s3.bucket, key, result.to_wire())
|
||||
except Exception as exc:
|
||||
logger.opt(exception=exc).error(
|
||||
"Не удалось записать файл результата. error_category=s3_error, job_id={job_id}, key={key}",
|
||||
job_id=request.job_id,
|
||||
key=key,
|
||||
)
|
||||
message = f"не удалось записать результат s3://{request.s3.bucket}/{key}: {_describe(exc)}"
|
||||
return self._error(request, ErrorCode.UNKNOWN_ERROR, message, outcome.metrics, started)
|
||||
|
||||
return self._envelope(
|
||||
job_id=request.job_id,
|
||||
file_id=request.file_id,
|
||||
status=MessageStatus(outcome.status.value),
|
||||
result_s3=S3Ref(bucket=request.s3.bucket, key=key),
|
||||
metrics=metrics,
|
||||
)
|
||||
|
||||
def _error(
|
||||
self, request: RecognitionRequest, code: ErrorCode, message: str, metrics: Metrics | None, started: float
|
||||
) -> RecognitionResultMessage:
|
||||
return self._envelope(
|
||||
job_id=request.job_id,
|
||||
file_id=request.file_id,
|
||||
status=MessageStatus.ERROR,
|
||||
error=MessageError(code=code, message=_one_line(message)),
|
||||
metrics=self._metrics(metrics, started),
|
||||
)
|
||||
|
||||
def _metrics(self, metrics: Metrics | None, started: float) -> Metrics:
|
||||
"""Метрики парсера + `duration_ms` (всегда, ≥ 0)."""
|
||||
duration_ms = max(0, round((self._clock() - started) * 1000))
|
||||
fields = metrics.model_dump(exclude_unset=True) if metrics is not None else {}
|
||||
return Metrics.model_validate({**fields, "duration_ms": duration_ms})
|
||||
|
||||
def _envelope( # noqa: PLR0913 - поля конверта именованные
|
||||
self,
|
||||
*,
|
||||
job_id: object,
|
||||
file_id: object,
|
||||
status: MessageStatus,
|
||||
error: MessageError | None = None,
|
||||
result_s3: S3Ref | None = None,
|
||||
metrics: Metrics | None = None,
|
||||
) -> RecognitionResultMessage:
|
||||
# Необязательные поля без null в схеме передаются, только если заданы (to_wire выводит заданные)
|
||||
fields: dict[str, Any] = {
|
||||
"schema_version": CONTRACT_VERSION,
|
||||
"job_id": job_id,
|
||||
"file_id": file_id,
|
||||
"status": status,
|
||||
"service": self._service,
|
||||
}
|
||||
if result_s3 is not None:
|
||||
fields["result_s3"] = result_s3
|
||||
fields["result_schema_version"] = CONTRACT_VERSION
|
||||
if error is not None:
|
||||
fields["error"] = error
|
||||
if metrics is not None:
|
||||
fields["metrics"] = metrics
|
||||
return RecognitionResultMessage.model_validate(fields)
|
||||
|
||||
async def _publish(self, envelope: RecognitionResultMessage, msg_id: str, job_id: str) -> None:
|
||||
"""Публикация конверта; исчерпанные повторы и прочие сбои — error-лог, не исключение (D-5)."""
|
||||
payload = json.dumps(envelope.to_wire(), ensure_ascii=False).encode("utf-8")
|
||||
try:
|
||||
await self._publisher.publish(RESULT_SUBJECT, payload, msg_id)
|
||||
except NatsPublishError as exc:
|
||||
logger.error(
|
||||
"Конверт не опубликован после всех повторов; ядро закроет задание по reply_deadline. "
|
||||
"error_category=nats_error, job_id={job_id}, msg_id={msg_id}, status={status}, error={error}",
|
||||
job_id=job_id,
|
||||
msg_id=msg_id,
|
||||
status=envelope.status.value,
|
||||
error=str(exc),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.opt(exception=exc).error(
|
||||
"Сбой публикации конверта. error_category=nats_error, job_id={job_id}, msg_id={msg_id}",
|
||||
job_id=job_id,
|
||||
msg_id=msg_id,
|
||||
)
|
||||
@ -38,6 +38,12 @@ class ParseOutcome:
|
||||
|
||||
@runtime_checkable
|
||||
class Parser(Protocol):
|
||||
"""Разбор файла: байты + тип → `ParseOutcome`.
|
||||
|
||||
Ожидаемые сбои разбора (не тот формат, нечитаемый лист и т. п.) возвращаются исходом `status = error` с
|
||||
`error {code, message}`; непредвиденное исключение воркер превращает в `UNKNOWN_ERROR` (design D-5).
|
||||
"""
|
||||
|
||||
def parse(self, data: bytes, file_type: FileType) -> ParseOutcome: ...
|
||||
|
||||
|
||||
|
||||
741
api/tests/unit/test_worker.py
Normal file
741
api/tests/unit/test_worker.py
Normal file
@ -0,0 +1,741 @@
|
||||
"""Unit-тесты воркера `src/core/worker.py` с фейками `Storage` / `Publisher` / `Parser` (design D-2, D-5, D-7).
|
||||
|
||||
Сценарии spec tabular-recognition-service: ack раньше результата, повторный запрос (во время обработки и после),
|
||||
заглушка `NOT_IMPLEMENTED`, `FILE_READ_ERROR`, `SCHEMA_MISMATCH` без ack, исключение парсера → `UNKNOWN_ERROR`,
|
||||
подмена парсера → `done` с файлом результата, повтор публикации и `NatsPublishError`. NATS и S3 не нужны.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from jsonschema import Draft202012Validator, FormatChecker
|
||||
from loguru import logger
|
||||
from referencing import Registry, Resource
|
||||
|
||||
from src.contract import (
|
||||
NATS_MSG_ID_HEADER,
|
||||
RESULT_SUBJECT,
|
||||
ErrorCode,
|
||||
FileType,
|
||||
MessageError,
|
||||
Metrics,
|
||||
OsvDocument,
|
||||
RecognitionResult,
|
||||
RecognitionResultMessage,
|
||||
ResultStatus,
|
||||
)
|
||||
from src.core.config import NATSSettings
|
||||
from src.core.nats import NatsConnection, NatsPublishError
|
||||
from src.core.worker import RecognitionWorker, result_key
|
||||
from src.parser import NullParser, ParseOutcome
|
||||
|
||||
API_DIR = Path(__file__).resolve().parents[2]
|
||||
CONTRACTS_DIR = API_DIR / "contracts"
|
||||
EXAMPLES_DIR = CONTRACTS_DIR / "examples"
|
||||
|
||||
REQUEST = json.loads((EXAMPLES_DIR / "request-tabular.json").read_text(encoding="utf-8"))
|
||||
JOB_ID = REQUEST["job_id"]
|
||||
FILE_ID = REQUEST["file_id"]
|
||||
BUCKET = REQUEST["s3"]["bucket"]
|
||||
SOURCE_KEY = REQUEST["s3"]["key"]
|
||||
RESULT_KEY = f"results/{FILE_ID}/{JOB_ID}/result.json"
|
||||
SERVICE = {"name": "tabular-parser", "version": "9.9.9"}
|
||||
OTHER_JOB_ID = "11111111-2222-4333-8444-555555555555"
|
||||
PARSER_ROWS = 3 # metrics.rows, которые возвращает подменённый парсер
|
||||
|
||||
|
||||
def _validators() -> dict[str, Draft202012Validator]:
|
||||
names = {
|
||||
"message": "recognition-result-message.schema.json",
|
||||
"result": "recognition-result.schema.json",
|
||||
}
|
||||
schemas = {kind: json.loads((CONTRACTS_DIR / name).read_text(encoding="utf-8")) for kind, name in names.items()}
|
||||
registry: Registry = Registry()
|
||||
for schema in schemas.values():
|
||||
registry = registry.with_resource(schema["$id"], Resource.from_contents(schema))
|
||||
return {
|
||||
kind: Draft202012Validator(schema, registry=registry, format_checker=FormatChecker())
|
||||
for kind, schema in schemas.items()
|
||||
}
|
||||
|
||||
|
||||
VALIDATORS = _validators()
|
||||
|
||||
|
||||
# --- фейки ------------------------------------------------------------------------------------------------------------
|
||||
class FakeStorage:
|
||||
"""`Storage` в памяти; пишет события в общий журнал `events`."""
|
||||
|
||||
def __init__(self, events: list[tuple[str, str]]) -> None:
|
||||
self.events = events
|
||||
self.objects: dict[tuple[str, str], bytes] = {}
|
||||
self.json_objects: dict[tuple[str, str], Any] = {}
|
||||
self.get_calls: list[tuple[str, str]] = []
|
||||
self.gate: asyncio.Event | None = None # задан — get_bytes ждёт его (обработка «в процессе»)
|
||||
self.get_started = asyncio.Event()
|
||||
|
||||
async def get_bytes(self, bucket: str, key: str) -> bytes:
|
||||
self.events.append(("s3_get", key))
|
||||
self.get_calls.append((bucket, key))
|
||||
self.get_started.set()
|
||||
if self.gate is not None:
|
||||
await self.gate.wait()
|
||||
try:
|
||||
return self.objects[(bucket, key)]
|
||||
except KeyError:
|
||||
raise FileNotFoundError(f"NoSuchKey: {key}") from None
|
||||
|
||||
async def put_json(self, bucket: str, key: str, obj: Any) -> None:
|
||||
self.events.append(("s3_put", key))
|
||||
json.dumps(obj) # объект обязан быть JSON-сериализуемым
|
||||
self.json_objects[(bucket, key)] = obj
|
||||
|
||||
|
||||
class FakePublisher:
|
||||
"""`Publisher`: копит опубликованное; `fail_subjects_msg_ids` — msg_id, публикация которых падает."""
|
||||
|
||||
def __init__(self, events: list[tuple[str, str]]) -> None:
|
||||
self.events = events
|
||||
self.published: list[tuple[str, dict[str, Any], str]] = []
|
||||
self.fail_msg_ids: set[str] = set()
|
||||
|
||||
async def publish(self, subject: str, payload: bytes, msg_id: str) -> None:
|
||||
if msg_id in self.fail_msg_ids:
|
||||
self.events.append(("publish_failed", msg_id))
|
||||
raise NatsPublishError(f"публикация {msg_id} не удалась")
|
||||
body = json.loads(payload)
|
||||
self.events.append(("publish", msg_id))
|
||||
self.published.append((subject, body, msg_id))
|
||||
|
||||
|
||||
class OsvParser:
|
||||
"""Подмена парсера (spec «Подмена реализации в тесте»): один `osvDocument` из примера контракта."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
example = json.loads((EXAMPLES_DIR / "result-osv-minimal.json").read_text(encoding="utf-8"))
|
||||
self.document = OsvDocument.model_validate(example["documents"][0])
|
||||
self.calls: list[tuple[bytes, FileType]] = []
|
||||
|
||||
def parse(self, data: bytes, file_type: FileType) -> ParseOutcome:
|
||||
self.calls.append((data, file_type))
|
||||
return ParseOutcome(status=ResultStatus.DONE, documents=(self.document,), metrics=Metrics(rows=PARSER_ROWS))
|
||||
|
||||
|
||||
class ExplodingParser:
|
||||
def parse(self, data: bytes, file_type: FileType) -> ParseOutcome:
|
||||
raise ValueError("лист без заголовка")
|
||||
|
||||
|
||||
class FakeClock:
|
||||
"""Монотонные часы: каждое обращение сдвигает время на `step` секунд."""
|
||||
|
||||
def __init__(self, step: float = 0.25) -> None:
|
||||
self.now = 100.0
|
||||
self.step = step
|
||||
|
||||
def __call__(self) -> float:
|
||||
self.now += self.step
|
||||
return self.now
|
||||
|
||||
|
||||
def _msg(body: Any, *, raw: bytes | None = None, msg_id: str | None = None) -> SimpleNamespace:
|
||||
data = raw if raw is not None else json.dumps(body).encode("utf-8")
|
||||
if msg_id is None and isinstance(body, dict) and "job_id" in body:
|
||||
msg_id = f"{body['job_id']}:{body.get('attempt')}"
|
||||
headers = {NATS_MSG_ID_HEADER: msg_id} if msg_id else None
|
||||
return SimpleNamespace(subject="recognition.tabular.request", data=data, headers=headers)
|
||||
|
||||
|
||||
def _request(**overrides: Any) -> dict[str, Any]:
|
||||
body = copy.deepcopy(REQUEST)
|
||||
body.update(overrides)
|
||||
return body
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def log_records() -> Iterator[list[Any]]:
|
||||
records: list[Any] = []
|
||||
handler_id = logger.add(lambda message: records.append(message.record), level="DEBUG")
|
||||
yield records
|
||||
logger.remove(handler_id)
|
||||
|
||||
|
||||
class Env:
|
||||
def __init__(self, parser: Any = None, *, concurrency: int = 1, done_cache_size: int = 1000) -> None:
|
||||
self.events: list[tuple[str, str]] = []
|
||||
self.storage = FakeStorage(self.events)
|
||||
self.storage.objects[(BUCKET, SOURCE_KEY)] = b"PK\x03\x04 xlsx bytes"
|
||||
self.publisher = FakePublisher(self.events)
|
||||
self.worker = RecognitionWorker(
|
||||
storage=self.storage,
|
||||
publisher=self.publisher,
|
||||
parser=parser if parser is not None else NullParser(),
|
||||
service_name=SERVICE["name"],
|
||||
service_version=SERVICE["version"],
|
||||
concurrency=concurrency,
|
||||
done_cache_size=done_cache_size,
|
||||
clock=FakeClock(),
|
||||
)
|
||||
|
||||
def bodies(self) -> list[dict[str, Any]]:
|
||||
return [body for _subject, body, _msg_id in self.publisher.published]
|
||||
|
||||
def msg_ids(self) -> list[str]:
|
||||
return [msg_id for _subject, _body, msg_id in self.publisher.published]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def env() -> Any:
|
||||
environment = Env()
|
||||
environment.worker.start()
|
||||
yield environment
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
|
||||
async def _make_env(parser: Any = None, **kwargs: Any) -> Env:
|
||||
environment = Env(parser, **kwargs)
|
||||
environment.worker.start()
|
||||
return environment
|
||||
|
||||
|
||||
def _assert_message_valid(body: dict[str, Any]) -> None:
|
||||
errors = [error.message for error in VALIDATORS["message"].iter_errors(body)]
|
||||
assert errors == [], errors
|
||||
RecognitionResultMessage.model_validate(body)
|
||||
|
||||
|
||||
# --- ack раньше результата ------------------------------------------------------------------------------------------
|
||||
async def test_ack_is_published_before_s3_read_and_result(env: Env) -> None:
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.events == [
|
||||
("publish", f"ack:{JOB_ID}:1"),
|
||||
("s3_get", SOURCE_KEY),
|
||||
("publish", f"result:{JOB_ID}:1"),
|
||||
]
|
||||
subjects = {subject for subject, _body, _msg_id in env.publisher.published}
|
||||
assert subjects == {RESULT_SUBJECT}
|
||||
ack = env.bodies()[0]
|
||||
assert ack == {
|
||||
"schema_version": "1.0",
|
||||
"job_id": JOB_ID,
|
||||
"file_id": FILE_ID,
|
||||
"status": "processing",
|
||||
"service": SERVICE,
|
||||
}
|
||||
for body in env.bodies():
|
||||
_assert_message_valid(body)
|
||||
|
||||
|
||||
async def test_ack_is_published_by_on_request_itself_before_queueing(env: Env) -> None:
|
||||
gate = asyncio.Event()
|
||||
env.storage.gate = gate
|
||||
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
|
||||
# on_request вернулся — ack уже опубликован, хотя обработка ещё не завершена
|
||||
assert env.msg_ids() == [f"ack:{JOB_ID}:1"]
|
||||
gate.set()
|
||||
await env.worker.join()
|
||||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||||
|
||||
|
||||
# --- заглушка парсера -----------------------------------------------------------------------------------------------
|
||||
async def test_null_parser_gives_not_implemented_with_duration(env: Env) -> None:
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
|
||||
result = env.bodies()[1]
|
||||
_assert_message_valid(result)
|
||||
assert result["status"] == "error"
|
||||
assert result["job_id"] == JOB_ID
|
||||
assert result["file_id"] == FILE_ID
|
||||
assert result["service"] == SERVICE
|
||||
assert result["error"] == {"code": "NOT_IMPLEMENTED", "message": "разбор XLS/XLSX/XLSB реализуется в Срезе 2"}
|
||||
assert "result_s3" not in result
|
||||
assert isinstance(result["metrics"]["duration_ms"], int)
|
||||
assert result["metrics"]["duration_ms"] >= 0
|
||||
assert env.storage.json_objects == {} # при ошибке файл результата не пишется
|
||||
|
||||
|
||||
# --- повторный запрос -----------------------------------------------------------------------------------------------
|
||||
async def test_repeated_request_during_processing_gets_second_ack_only(env: Env) -> None:
|
||||
gate = asyncio.Event()
|
||||
env.storage.gate = gate
|
||||
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await asyncio.wait_for(env.storage.get_started.wait(), 1)
|
||||
await env.worker.on_request(_msg(_request())) # та же доставка ещё раз, пока первая в работе
|
||||
gate.set()
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||||
assert env.storage.get_calls == [(BUCKET, SOURCE_KEY)]
|
||||
|
||||
|
||||
async def test_repeated_request_while_queued_is_not_processed_twice() -> None:
|
||||
environment = Env() # воркер не запущен: задание стоит в очереди
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
environment.worker.start()
|
||||
await environment.worker.join()
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
assert environment.msg_ids() == [f"ack:{JOB_ID}:1", f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||||
assert environment.storage.get_calls == [(BUCKET, SOURCE_KEY)]
|
||||
|
||||
|
||||
async def test_repeated_request_after_done_gets_ack_with_its_own_attempt(env: Env) -> None:
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
await env.worker.on_request(_msg(_request(attempt=2)))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1", f"ack:{JOB_ID}:2"]
|
||||
assert env.storage.get_calls == [(BUCKET, SOURCE_KEY)]
|
||||
|
||||
|
||||
async def test_failed_job_counts_as_processed(env: Env) -> None:
|
||||
del env.storage.objects[(BUCKET, SOURCE_KEY)]
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1", f"ack:{JOB_ID}:1"]
|
||||
assert len(env.storage.get_calls) == 1
|
||||
|
||||
|
||||
async def test_done_cache_is_bounded_lru() -> None:
|
||||
environment = await _make_env(done_cache_size=1)
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
await environment.worker.on_request(_msg(_request(job_id=OTHER_JOB_ID)))
|
||||
await environment.worker.join()
|
||||
# первый job_id вытеснен из LRU размера 1 — его повтор обрабатывается снова
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
assert len(environment.storage.get_calls) == 3 # noqa: PLR2004 - три запроса: первый, второй, повтор первого
|
||||
|
||||
|
||||
# --- FILE_READ_ERROR ------------------------------------------------------------------------------------------------
|
||||
async def test_missing_s3_object_gives_file_read_error_and_worker_continues(env: Env) -> None:
|
||||
del env.storage.objects[(BUCKET, SOURCE_KEY)]
|
||||
other_key = f"uploads/{FILE_ID}/{OTHER_JOB_ID}/osv.xlsx"
|
||||
env.storage.objects[(BUCKET, other_key)] = b"bytes"
|
||||
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
await env.worker.on_request(_msg(_request(job_id=OTHER_JOB_ID, s3={"bucket": BUCKET, "key": other_key})))
|
||||
await env.worker.join()
|
||||
|
||||
first, second = env.bodies()[1], env.bodies()[3]
|
||||
_assert_message_valid(first)
|
||||
assert first["status"] == "error"
|
||||
assert first["error"]["code"] == "FILE_READ_ERROR"
|
||||
assert SOURCE_KEY in first["error"]["message"]
|
||||
assert "result_s3" not in first
|
||||
assert first["metrics"]["duration_ms"] >= 0
|
||||
# следующий запрос обработан как обычно
|
||||
assert second["job_id"] == OTHER_JOB_ID
|
||||
assert second["error"]["code"] == "NOT_IMPLEMENTED"
|
||||
|
||||
|
||||
# --- SCHEMA_MISMATCH ------------------------------------------------------------------------------------------------
|
||||
async def test_unsupported_major_version_gives_schema_mismatch_without_ack(env: Env) -> None:
|
||||
await env.worker.on_request(_msg(_request(schema_version="2.0")))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"result:{JOB_ID}:1"]
|
||||
result = env.bodies()[0]
|
||||
_assert_message_valid(result)
|
||||
assert result["status"] == "error"
|
||||
assert result["schema_version"] == "1.0"
|
||||
assert result["job_id"] == JOB_ID
|
||||
assert result["file_id"] == FILE_ID
|
||||
assert result["error"]["code"] == "SCHEMA_MISMATCH"
|
||||
assert "2.0" in result["error"]["message"]
|
||||
assert result["metrics"]["duration_ms"] >= 0
|
||||
assert env.storage.get_calls == []
|
||||
|
||||
|
||||
async def test_minor_version_is_accepted(env: Env) -> None:
|
||||
await env.worker.on_request(_msg(_request(schema_version="1.7", extra_field="ignored")))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mutate",
|
||||
[
|
||||
pytest.param(lambda body: body["s3"].pop("key"), id="no-s3-key"),
|
||||
pytest.param(lambda body: body.update(file_type="DOCX"), id="unknown-file-type"),
|
||||
pytest.param(lambda body: body.pop("reply_deadline"), id="no-reply-deadline"),
|
||||
],
|
||||
)
|
||||
async def test_request_violating_schema_gives_schema_mismatch_without_ack(env: Env, mutate: Any) -> None:
|
||||
body = _request()
|
||||
mutate(body)
|
||||
|
||||
await env.worker.on_request(_msg(body))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"result:{JOB_ID}:1"]
|
||||
result = env.bodies()[0]
|
||||
_assert_message_valid(result)
|
||||
assert result["error"]["code"] == "SCHEMA_MISMATCH"
|
||||
assert env.storage.get_calls == []
|
||||
|
||||
|
||||
async def test_schema_mismatch_takes_attempt_from_header_when_body_attempt_is_invalid(env: Env) -> None:
|
||||
await env.worker.on_request(_msg(_request(attempt=0), msg_id=f"{JOB_ID}:3"))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"result:{JOB_ID}:3"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"message",
|
||||
[
|
||||
pytest.param(_msg(None, raw=b"\xff not json"), id="not-json"),
|
||||
pytest.param(_msg(None, raw=b"[1, 2]"), id="json-array"),
|
||||
pytest.param(_msg(_request(job_id="not-a-uuid")), id="job-id-unreadable"),
|
||||
pytest.param(_msg({k: v for k, v in REQUEST.items() if k != "file_id"}), id="no-file-id"),
|
||||
pytest.param(_msg(_request(attempt="x"), msg_id="garbage"), id="attempt-unreadable"),
|
||||
],
|
||||
)
|
||||
async def test_unreadable_request_is_logged_and_dropped(env: Env, log_records: list[Any], message: Any) -> None:
|
||||
await env.worker.on_request(message)
|
||||
await env.worker.join()
|
||||
|
||||
assert env.publisher.published == []
|
||||
assert env.storage.get_calls == []
|
||||
errors = [record for record in log_records if record["level"].name == "ERROR"]
|
||||
assert errors, "отброшенный запрос должен попасть в лог уровня error"
|
||||
assert "error_category=schema_mismatch" in errors[0]["message"]
|
||||
|
||||
|
||||
# --- исключение парсера ---------------------------------------------------------------------------------------------
|
||||
async def test_parser_exception_gives_unknown_error_without_traceback() -> None:
|
||||
environment = await _make_env(ExplodingParser())
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
result = environment.bodies()[1]
|
||||
_assert_message_valid(result)
|
||||
assert result["error"]["code"] == "UNKNOWN_ERROR"
|
||||
assert "лист без заголовка" in result["error"]["message"]
|
||||
assert "Traceback" not in result["error"]["message"]
|
||||
assert "\n" not in result["error"]["message"]
|
||||
assert result["metrics"]["duration_ms"] >= 0
|
||||
|
||||
|
||||
async def test_parser_is_called_in_a_thread(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls: list[Any] = []
|
||||
original = asyncio.to_thread
|
||||
|
||||
async def spy(func: Any, /, *args: Any, **kwargs: Any) -> Any:
|
||||
calls.append(func)
|
||||
return await original(func, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(asyncio, "to_thread", spy)
|
||||
parser = OsvParser()
|
||||
environment = await _make_env(parser)
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
assert calls == [parser.parse]
|
||||
assert parser.calls == [(b"PK\x03\x04 xlsx bytes", "XLSX")]
|
||||
|
||||
|
||||
# --- подмена парсера → done -----------------------------------------------------------------------------------------
|
||||
async def test_substituted_parser_gives_done_with_result_file() -> None:
|
||||
parser = OsvParser()
|
||||
environment = await _make_env(parser)
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
assert environment.events == [
|
||||
("publish", f"ack:{JOB_ID}:1"),
|
||||
("s3_get", SOURCE_KEY),
|
||||
("s3_put", RESULT_KEY),
|
||||
("publish", f"result:{JOB_ID}:1"),
|
||||
]
|
||||
result_file = environment.storage.json_objects[(BUCKET, RESULT_KEY)]
|
||||
errors = [error.message for error in VALIDATORS["result"].iter_errors(result_file)]
|
||||
assert errors == [], errors
|
||||
parsed = RecognitionResult.model_validate(result_file)
|
||||
assert parsed.status is ResultStatus.DONE
|
||||
assert parsed.documents == [parser.document]
|
||||
assert str(parsed.job_id) == JOB_ID
|
||||
assert result_file["service"] == SERVICE
|
||||
assert result_file["metrics"]["rows"] == PARSER_ROWS
|
||||
|
||||
envelope = environment.bodies()[1]
|
||||
_assert_message_valid(envelope)
|
||||
assert envelope["status"] == "done"
|
||||
assert envelope["result_s3"] == {"bucket": BUCKET, "key": RESULT_KEY}
|
||||
assert envelope["result_schema_version"] == "1.0"
|
||||
assert "error" not in envelope
|
||||
assert envelope["metrics"]["rows"] == PARSER_ROWS
|
||||
assert envelope["metrics"]["duration_ms"] >= 0
|
||||
|
||||
|
||||
async def test_needs_check_outcome_is_propagated() -> None:
|
||||
parser = OsvParser()
|
||||
|
||||
class NeedsCheckParser:
|
||||
def parse(self, data: bytes, file_type: FileType) -> ParseOutcome:
|
||||
return ParseOutcome(status=ResultStatus.NEEDS_CHECK, documents=(parser.document,))
|
||||
|
||||
environment = await _make_env(NeedsCheckParser())
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
envelope = environment.bodies()[1]
|
||||
_assert_message_valid(envelope)
|
||||
assert envelope["status"] == "needs_check"
|
||||
assert environment.storage.json_objects[(BUCKET, RESULT_KEY)]["status"] == "needs_check"
|
||||
|
||||
|
||||
async def test_error_outcome_metrics_are_kept_and_duration_added() -> None:
|
||||
class ErrorWithMetricsParser:
|
||||
def parse(self, data: bytes, file_type: FileType) -> ParseOutcome:
|
||||
return ParseOutcome(
|
||||
status=ResultStatus.ERROR,
|
||||
error=MessageError(code=ErrorCode.INVALID_FORMAT, message="не ОСВ"),
|
||||
metrics=Metrics(warnings=["лист пуст"]),
|
||||
)
|
||||
|
||||
environment = await _make_env(ErrorWithMetricsParser())
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
envelope = environment.bodies()[1]
|
||||
_assert_message_valid(envelope)
|
||||
assert envelope["error"] == {"code": "INVALID_FORMAT", "message": "не ОСВ"}
|
||||
assert envelope["metrics"]["warnings"] == ["лист пуст"]
|
||||
assert envelope["metrics"]["duration_ms"] >= 0
|
||||
|
||||
|
||||
async def test_result_file_write_failure_gives_error_without_result_s3() -> None:
|
||||
environment = await _make_env(OsvParser())
|
||||
|
||||
async def broken_put(bucket: str, key: str, obj: Any) -> None:
|
||||
raise ConnectionError("S3 недоступен")
|
||||
|
||||
environment.storage.put_json = broken_put # type: ignore[method-assign]
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
envelope = environment.bodies()[1]
|
||||
_assert_message_valid(envelope)
|
||||
assert envelope["status"] == "error"
|
||||
assert envelope["error"]["code"] == "UNKNOWN_ERROR"
|
||||
assert "result_s3" not in envelope
|
||||
|
||||
|
||||
def test_result_key_template() -> None:
|
||||
assert result_key(FILE_ID, JOB_ID) == RESULT_KEY
|
||||
|
||||
|
||||
# --- публикация: повтор и исчерпание попыток ------------------------------------------------------------------------
|
||||
async def test_result_publish_failure_is_logged_and_worker_continues(env: Env, log_records: list[Any]) -> None:
|
||||
env.publisher.fail_msg_ids.add(f"result:{JOB_ID}:1")
|
||||
env.storage.objects[(BUCKET, "uploads/other.xlsx")] = b"bytes"
|
||||
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
await env.worker.on_request(_msg(_request(job_id=OTHER_JOB_ID, s3={"bucket": BUCKET, "key": "uploads/other.xlsx"})))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"ack:{OTHER_JOB_ID}:1", f"result:{OTHER_JOB_ID}:1"]
|
||||
errors = [record["message"] for record in log_records if record["level"].name == "ERROR"]
|
||||
assert any(f"result:{JOB_ID}:1" in message for message in errors)
|
||||
|
||||
|
||||
async def test_ack_publish_failure_is_logged_and_job_still_processed(env: Env, log_records: list[Any]) -> None:
|
||||
env.publisher.fail_msg_ids.add(f"ack:{JOB_ID}:1")
|
||||
|
||||
await env.worker.on_request(_msg(_request()))
|
||||
await env.worker.join()
|
||||
|
||||
assert env.msg_ids() == [f"result:{JOB_ID}:1"]
|
||||
assert env.storage.get_calls == [(BUCKET, SOURCE_KEY)]
|
||||
errors = [record["message"] for record in log_records if record["level"].name == "ERROR"]
|
||||
assert any(f"ack:{JOB_ID}:1" in message for message in errors)
|
||||
|
||||
|
||||
class _FlakyNatsClient:
|
||||
"""Клиент nats-py для `NatsConnection`: первая публикация результата падает, остальные проходят."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.is_connected = False
|
||||
self.is_closed = False
|
||||
self.attempts: list[tuple[str, dict[str, str]]] = []
|
||||
self.delivered: list[tuple[str, dict[str, str]]] = []
|
||||
self._failed_once = False
|
||||
|
||||
async def connect(self, *_args: Any, **_kwargs: Any) -> None:
|
||||
self.is_connected = True
|
||||
|
||||
async def publish(self, subject: str, payload: bytes, headers: dict[str, str]) -> None:
|
||||
msg_id = headers[NATS_MSG_ID_HEADER]
|
||||
self.attempts.append((msg_id, headers))
|
||||
if msg_id.startswith("result:") and not self._failed_once:
|
||||
self._failed_once = True
|
||||
raise ConnectionError("соединение сброшено")
|
||||
self.delivered.append((msg_id, json.loads(payload)))
|
||||
|
||||
async def flush(self, timeout: float) -> 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
|
||||
|
||||
|
||||
async def test_result_publish_is_retried_with_same_msg_id_through_nats_connection() -> None:
|
||||
client = _FlakyNatsClient()
|
||||
pauses: list[float] = []
|
||||
|
||||
async def fake_sleep(seconds: float) -> None:
|
||||
pauses.append(seconds)
|
||||
|
||||
connection = NatsConnection(
|
||||
NATSSettings(NATS_URL="nats://fake:4222"),
|
||||
client_factory=lambda: client,
|
||||
sleep=fake_sleep,
|
||||
publish_retry_pause=2.0,
|
||||
)
|
||||
await connection.start()
|
||||
assert await connection.wait_ready(1)
|
||||
events: list[tuple[str, str]] = []
|
||||
storage = FakeStorage(events)
|
||||
storage.objects[(BUCKET, SOURCE_KEY)] = b"bytes"
|
||||
worker = RecognitionWorker(
|
||||
storage=storage,
|
||||
publisher=connection,
|
||||
parser=NullParser(),
|
||||
service_name=SERVICE["name"],
|
||||
service_version=SERVICE["version"],
|
||||
)
|
||||
worker.start()
|
||||
try:
|
||||
await worker.on_request(_msg(_request()))
|
||||
await worker.join()
|
||||
finally:
|
||||
await worker.stop(timeout=1)
|
||||
await connection.stop()
|
||||
|
||||
assert [msg_id for msg_id, _headers in client.attempts] == [
|
||||
f"ack:{JOB_ID}:1",
|
||||
f"result:{JOB_ID}:1",
|
||||
f"result:{JOB_ID}:1",
|
||||
]
|
||||
assert [msg_id for msg_id, _body in client.delivered] == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||||
assert pauses == [2.0]
|
||||
assert storage.get_calls == [(BUCKET, SOURCE_KEY)] # обработка не повторялась
|
||||
|
||||
|
||||
# --- конкуренция и остановка ----------------------------------------------------------------------------------------
|
||||
async def test_concurrency_processes_jobs_in_parallel() -> None:
|
||||
JOBS = 2 # noqa: N806 - число одновременных заданий
|
||||
environment = await _make_env(concurrency=2)
|
||||
gate = asyncio.Event()
|
||||
environment.storage.gate = gate
|
||||
started: list[str] = []
|
||||
original_get = environment.storage.get_bytes
|
||||
|
||||
async def tracking_get(bucket: str, key: str) -> bytes:
|
||||
started.append(key)
|
||||
return await original_get(bucket, key)
|
||||
|
||||
environment.storage.get_bytes = tracking_get # type: ignore[method-assign]
|
||||
other_key = "uploads/other.xlsx"
|
||||
environment.storage.objects[(BUCKET, other_key)] = b"bytes"
|
||||
try:
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.on_request(
|
||||
_msg(_request(job_id=OTHER_JOB_ID, s3={"bucket": BUCKET, "key": other_key}))
|
||||
)
|
||||
for _ in range(50):
|
||||
if len(started) == JOBS:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert len(started) == JOBS # оба задания одновременно ждут S3
|
||||
gate.set()
|
||||
await environment.worker.join()
|
||||
finally:
|
||||
await environment.worker.stop(timeout=1)
|
||||
|
||||
|
||||
async def test_stop_waits_for_in_flight_job_then_rejects_new_requests() -> None:
|
||||
environment = await _make_env()
|
||||
gate = asyncio.Event()
|
||||
environment.storage.gate = gate
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await asyncio.wait_for(environment.storage.get_started.wait(), 1)
|
||||
|
||||
stop_task = asyncio.create_task(environment.worker.stop(timeout=5))
|
||||
await asyncio.sleep(0)
|
||||
# после начала остановки новые запросы не принимаются (без ack — ядро переотправит другой инстанции)
|
||||
await environment.worker.on_request(_msg(_request(job_id=OTHER_JOB_ID)))
|
||||
gate.set()
|
||||
await asyncio.wait_for(stop_task, 2)
|
||||
|
||||
assert environment.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||||
|
||||
|
||||
async def test_stop_cancels_jobs_that_do_not_finish_in_time(log_records: list[Any]) -> None:
|
||||
environment = await _make_env()
|
||||
environment.storage.gate = asyncio.Event() # никогда не откроется
|
||||
await environment.worker.on_request(_msg(_request()))
|
||||
await environment.worker.on_request(_msg(_request(job_id=OTHER_JOB_ID)))
|
||||
await asyncio.wait_for(environment.storage.get_started.wait(), 1)
|
||||
|
||||
await asyncio.wait_for(environment.worker.stop(timeout=0.05), 2)
|
||||
|
||||
assert environment.msg_ids() == [f"ack:{JOB_ID}:1", f"ack:{OTHER_JOB_ID}:1"]
|
||||
warnings = [record["message"] for record in log_records if record["level"].name == "WARNING"]
|
||||
assert any("pending=1" in message for message in warnings)
|
||||
|
||||
|
||||
async def test_stop_is_idempotent_and_safe_without_start() -> None:
|
||||
environment = Env()
|
||||
await environment.worker.stop(timeout=0.1)
|
||||
await environment.worker.stop(timeout=0.1)
|
||||
Loading…
x
Reference in New Issue
Block a user