3a08b990dd
Refs #14 Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com> Claude-Session: cedd7915-2efe-41e5-988c-a9117033598b
857 lines
35 KiB
Python
857 lines
35 KiB
Python
"""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, get_parser
|
||
|
||
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_msg_ids` — msg_id, публикация которых падает; `gates` — msg_id,
|
||
публикация которых ждёт события (`waiting` выставляется, когда публикация встала на ожидание)."""
|
||
|
||
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()
|
||
self.gates: dict[str, asyncio.Event] = {}
|
||
self.waiting = asyncio.Event()
|
||
|
||
async def publish(self, subject: str, payload: bytes, msg_id: str) -> None:
|
||
gate = self.gates.get(msg_id)
|
||
if gate is not None:
|
||
self.waiting.set()
|
||
await gate.wait()
|
||
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.2",
|
||
"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.2"
|
||
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.2"
|
||
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] = []
|
||
all_started = asyncio.Event()
|
||
original_get = environment.storage.get_bytes
|
||
|
||
async def tracking_get(bucket: str, key: str) -> bytes:
|
||
started.append(key)
|
||
if len(started) == JOBS:
|
||
all_started.set()
|
||
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}))
|
||
)
|
||
# оба задания одновременно ждут S3 (первое не отпущено): при concurrency=1 второе не стартовало бы
|
||
await asyncio.wait_for(all_started.wait(), 1)
|
||
assert len(started) == JOBS
|
||
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)
|
||
|
||
|
||
async def test_stop_during_ack_publish_drops_job_with_warning(log_records: list[Any]) -> None:
|
||
environment = await _make_env()
|
||
gate = asyncio.Event()
|
||
environment.publisher.gates[f"ack:{JOB_ID}:1"] = gate
|
||
request_task = asyncio.create_task(environment.worker.on_request(_msg(_request())))
|
||
await asyncio.wait_for(environment.publisher.waiting.wait(), 1)
|
||
|
||
# stop() начался, пока on_request публикует ack: очередь пуста, корутины отменены
|
||
await asyncio.wait_for(environment.worker.stop(timeout=1), 2)
|
||
gate.set()
|
||
await asyncio.wait_for(request_task, 1)
|
||
|
||
assert environment.msg_ids() == [f"ack:{JOB_ID}:1"]
|
||
assert environment.storage.get_calls == []
|
||
assert environment.worker._in_progress == set()
|
||
assert environment.worker._queue.qsize() == 0
|
||
warnings = [record["message"] for record in log_records if record["level"].name == "WARNING"]
|
||
assert any("error_category=shutdown" in message and JOB_ID in message for message in warnings), warnings
|
||
|
||
|
||
# --- нормализация сообщения ошибки парсера ------------------------------------------------------------------------
|
||
async def test_parser_error_message_is_normalised_to_one_bounded_line() -> None:
|
||
long_message = "строка 1\nстрока 2\t" + "x" * 3000
|
||
|
||
class VerboseErrorParser:
|
||
def parse(self, data: bytes, file_type: FileType) -> ParseOutcome:
|
||
return ParseOutcome(
|
||
status=ResultStatus.ERROR,
|
||
error=MessageError(code=ErrorCode.INVALID_FORMAT, message=long_message),
|
||
)
|
||
|
||
environment = await _make_env(VerboseErrorParser())
|
||
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)
|
||
message = envelope["error"]["message"]
|
||
assert envelope["error"]["code"] == "INVALID_FORMAT"
|
||
assert message.startswith("строка 1 строка 2 xxx")
|
||
assert "\n" not in message
|
||
assert "\t" not in message
|
||
assert len(message) <= 1000 # noqa: PLR2004 - ERROR_MESSAGE_LIMIT
|
||
|
||
|
||
# --- невалидное тело с job_id, который уже в работе или обработан ---------------------------------------------------
|
||
async def test_malformed_request_for_job_in_progress_is_dropped(env: Env, log_records: list[Any]) -> 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)
|
||
|
||
malformed = _request()
|
||
malformed["s3"].pop("key")
|
||
await env.worker.on_request(_msg(malformed))
|
||
gate.set()
|
||
await env.worker.join()
|
||
|
||
# SCHEMA_MISMATCH под тем же result:<job_id>:<attempt> не публикуется — только настоящий результат
|
||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||
assert env.bodies()[1]["error"]["code"] == "NOT_IMPLEMENTED"
|
||
warnings = [record["message"] for record in log_records if record["level"].name == "WARNING"]
|
||
assert any("error_category=schema_mismatch" in message and JOB_ID in message for message in warnings)
|
||
|
||
|
||
async def test_malformed_request_for_done_job_is_dropped(env: Env) -> None:
|
||
await env.worker.on_request(_msg(_request()))
|
||
await env.worker.join()
|
||
|
||
await env.worker.on_request(_msg(_request(job_id=JOB_ID.upper(), schema_version="2.0")))
|
||
await env.worker.join()
|
||
|
||
assert env.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||
|
||
|
||
# --- реальный парсер в воркере ------------------------------------------------------------------------------------
|
||
REAL_OSV_66 = Path("osv") / "rb" / "2024-FY" / "ОСВ_66 за 2024.xls"
|
||
|
||
|
||
@pytest.mark.fixtures
|
||
async def test_worker_with_real_osv_parser_gives_done_and_valid_result_file(fixtures_dir: Path) -> None:
|
||
"""Spec «Реальная ОСВ через воркер»: настоящий `OsvParser` (`get_parser`) + байты реальной ОСВ 66 без NATS/S3."""
|
||
environment = Env(get_parser())
|
||
environment.storage.objects[(BUCKET, SOURCE_KEY)] = (fixtures_dir / REAL_OSV_66).read_bytes()
|
||
environment.worker.start()
|
||
try:
|
||
await environment.worker.on_request(_msg(_request(file_type="XLS")))
|
||
await environment.worker.join()
|
||
finally:
|
||
await environment.worker.stop(timeout=1)
|
||
|
||
assert environment.msg_ids() == [f"ack:{JOB_ID}:1", f"result:{JOB_ID}:1"]
|
||
envelope = environment.bodies()[1]
|
||
_assert_message_valid(envelope)
|
||
assert envelope["status"] == "done"
|
||
assert envelope["result_schema_version"] == "1.2"
|
||
assert envelope["result_s3"] == {"bucket": BUCKET, "key": RESULT_KEY}
|
||
stored = environment.storage.json_objects[(BUCKET, RESULT_KEY)]
|
||
assert [error.message for error in VALIDATORS["result"].iter_errors(stored)] == []
|
||
assert RecognitionResult.model_validate(stored).to_wire() == stored
|
||
assert [document["document_type"] for document in stored["documents"]] == ["OSV"]
|
||
assert stored["documents"][0]["header"]["account_code_raw"] == "66"
|