9b81fa8e37
handle_message(msg, *, session_factory, mapper, now): разбор конверта по контракту v1.0, одна короткая транзакция с SELECT … FOR UPDATE (ждёт блокировку публикатора/вахтёра); ack published → processing; результат — дедуп по result_msg_id, финальные статусы (R-22) не трогаются, регистрация s3_object (results), слияние метрик с сохранением republish_count. Подписка app.py получает handler с фабрикой сессий и NoopMapper; шаблонный core/nats_callback.py удалён (R-3). audit_event не пишется (R-7). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
202 lines
9.2 KiB
Python
202 lines
9.2 KiB
Python
"""Фабрики тестовых данных для БД-тестов.
|
||
|
||
- `make_osv_entry` — минимальная цепочка FK поверх сессии (`flush`, без `commit`): откатывается фикстурой `db_session`.
|
||
- `CommittedJobs` — задания распознавания, ЗАФИКСИРОВАННЫЕ в БД (`commit`): для фоновых процессов обмена, которые
|
||
работают в собственных транзакциях и не видят чужую незафиксированную транзакцию. Всё созданное удаляется
|
||
`cleanup()` в порядке FK, включая строки `s3_object` файлов результата, зарегистрированные inbox'ом.
|
||
"""
|
||
|
||
import datetime
|
||
import decimal
|
||
import uuid
|
||
from typing import Any
|
||
|
||
from sqlalchemy import delete, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||
|
||
from src.models.directories import LegalEntity
|
||
from src.models.enums import (
|
||
DocumentStatus,
|
||
DocumentType,
|
||
FileStatus,
|
||
FileType,
|
||
RecognitionPipeline,
|
||
RecognitionStatus,
|
||
S3BucketRole,
|
||
)
|
||
from src.models.files import Document, RecognitionJob, S3Object, StoredFile
|
||
from src.models.osv import OsvEntry
|
||
|
||
DEFAULT_FILE_TYPES = {RecognitionPipeline.tabular: FileType.XLSX, RecognitionPipeline.image: FileType.PDF}
|
||
|
||
|
||
async def make_osv_entry(session: AsyncSession, **overrides: Any) -> OsvEntry:
|
||
"""Строка ОСВ с минимальной цепочкой FK (ЮЛ, s3_object, stored_file, recognition_job, document); `flush`.
|
||
|
||
ИНН ЮЛ случайный, поэтому фабрику можно вызывать несколько раз в одном тесте.
|
||
|
||
`overrides` перекрывают поля `OsvEntry` (например, `closing_debit=Decimal("1.00")`). Данные остаются в
|
||
транзакции сессии: фикстура `db_session` откатывает их после теста.
|
||
"""
|
||
now = datetime.datetime.now(datetime.UTC)
|
||
entity = LegalEntity(inn=str(uuid.uuid4().int)[:10], name="ООО Тест", valid_from=now)
|
||
s3 = S3Object(bucket_role=S3BucketRole.uploads, bucket="uploads", key=f"{uuid.uuid4()}/osv.xlsx", created_at=now)
|
||
session.add_all([entity, s3])
|
||
await session.flush()
|
||
file = StoredFile(
|
||
legal_entity_id=entity.id, file_type=FileType.XLSX, s3_object_id=s3.id, status=FileStatus.UPLOADED
|
||
)
|
||
session.add(file)
|
||
await session.flush()
|
||
job = RecognitionJob(
|
||
file_id=file.id,
|
||
pipeline=RecognitionPipeline.tabular,
|
||
request_subject="recognition.tabular.request",
|
||
schema_version="1",
|
||
status=RecognitionStatus.queued,
|
||
)
|
||
session.add(job)
|
||
await session.flush()
|
||
document = Document(
|
||
stored_file_id=file.id,
|
||
legal_entity_id=entity.id,
|
||
document_type=DocumentType.OSV,
|
||
osv_account_code="62",
|
||
status=DocumentStatus.RECOGNIZED,
|
||
)
|
||
session.add(document)
|
||
await session.flush()
|
||
fields: dict[str, Any] = {
|
||
"recognition_job_id": job.id,
|
||
"document_id": document.id,
|
||
"owner_legal_entity_id": entity.id,
|
||
"row_number": 1,
|
||
"closing_debit": decimal.Decimal("0.00"),
|
||
}
|
||
fields.update(overrides)
|
||
entry = OsvEntry(**fields)
|
||
session.add(entry)
|
||
await session.flush()
|
||
return entry
|
||
|
||
|
||
class CommittedJobs:
|
||
"""Задания `recognition_job` с цепочкой `legal_entity` + `s3_object` + `stored_file`, зафиксированные в БД.
|
||
|
||
Каждый `make_queued_job` — отдельная транзакция с `commit`; созданные id запоминаются, `cleanup()` удаляет их
|
||
в порядке FK (задание → файл → объект S3 и ЮЛ). Вызывать `cleanup()` в teardown фикстуры — даже при падении теста.
|
||
"""
|
||
|
||
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
||
self._session_factory = session_factory
|
||
self.job_ids: list[uuid.UUID] = []
|
||
self._file_ids: list[uuid.UUID] = []
|
||
self._s3_ids: list[uuid.UUID] = []
|
||
self._extra_s3_ids: list[uuid.UUID] = []
|
||
self._entity_ids: list[uuid.UUID] = []
|
||
|
||
async def make_queued_job( # noqa: PLR0913 - каждое поле задания задаётся тестом явно
|
||
self,
|
||
*,
|
||
pipeline: RecognitionPipeline = RecognitionPipeline.tabular,
|
||
file_type: FileType | None = None,
|
||
original_name: str | None = "osv.xlsx",
|
||
status: RecognitionStatus = RecognitionStatus.queued,
|
||
reply_deadline: datetime.datetime | None = None,
|
||
published_at: datetime.datetime | None = None,
|
||
request_msg_id: str | None = None,
|
||
metrics: Any = None,
|
||
attempt: int = 1,
|
||
result_msg_id: str | None = None,
|
||
started_at: datetime.datetime | None = None,
|
||
) -> uuid.UUID:
|
||
"""Задание (по умолчанию `queued`, tabular, XLSX) со своей цепочкой FK; возвращает `recognition_job.id`."""
|
||
now = datetime.datetime.now(datetime.UTC)
|
||
async with self._session_factory() as session, session.begin():
|
||
entity = LegalEntity(inn=str(uuid.uuid4().int)[:12], name="ООО Тест обмена", valid_from=now)
|
||
s3 = S3Object(
|
||
bucket_role=S3BucketRole.uploads,
|
||
bucket="arm-ka",
|
||
key=f"uploads/{uuid.uuid4()}/{original_name or 'file'}",
|
||
created_at=now,
|
||
)
|
||
session.add_all([entity, s3])
|
||
await session.flush()
|
||
self._entity_ids.append(entity.id)
|
||
self._s3_ids.append(s3.id)
|
||
file = StoredFile(
|
||
legal_entity_id=entity.id,
|
||
file_type=file_type or DEFAULT_FILE_TYPES[pipeline],
|
||
original_name=original_name,
|
||
s3_object_id=s3.id,
|
||
status=FileStatus.UPLOADED,
|
||
)
|
||
session.add(file)
|
||
await session.flush()
|
||
self._file_ids.append(file.id)
|
||
job = RecognitionJob(
|
||
file_id=file.id,
|
||
pipeline=pipeline,
|
||
request_subject=f"recognition.{pipeline.value}.request",
|
||
schema_version="1.0",
|
||
status=status,
|
||
attempt=attempt,
|
||
reply_deadline=reply_deadline,
|
||
published_at=published_at,
|
||
request_msg_id=request_msg_id,
|
||
metrics=metrics,
|
||
result_msg_id=result_msg_id,
|
||
started_at=started_at,
|
||
)
|
||
session.add(job)
|
||
await session.flush()
|
||
self.job_ids.append(job.id)
|
||
return job.id
|
||
|
||
async def s3_key(self, job_id: uuid.UUID) -> str:
|
||
"""Ключ S3 исходного файла задания."""
|
||
async with self._session_factory() as session:
|
||
job = await session.get(RecognitionJob, job_id)
|
||
assert job is not None
|
||
file = await session.get(StoredFile, job.file_id)
|
||
assert file is not None
|
||
s3 = await session.get(S3Object, file.s3_object_id)
|
||
assert s3 is not None
|
||
return s3.key
|
||
|
||
async def get(self, job_id: uuid.UUID) -> RecognitionJob:
|
||
"""Свежее состояние задания из БД (новая сессия)."""
|
||
async with self._session_factory() as session:
|
||
job = await session.get(RecognitionJob, job_id)
|
||
assert job is not None
|
||
return job
|
||
|
||
def track_s3_object(self, s3_object_id: uuid.UUID) -> None:
|
||
"""Запомнить строку `s3_object`, созданную тестом вне фабрики (удаляется в `cleanup()`)."""
|
||
self._extra_s3_ids.append(s3_object_id)
|
||
|
||
async def cleanup(self) -> None:
|
||
"""Удаляет всё созданное в порядке FK (файлы результата, на которые ссылаются задания, — после заданий)."""
|
||
async with self._session_factory() as session, session.begin():
|
||
result_s3_ids: list[uuid.UUID] = []
|
||
if self.job_ids:
|
||
rows = await session.execute(
|
||
select(RecognitionJob.result_s3_object_id).where(
|
||
RecognitionJob.id.in_(self.job_ids), RecognitionJob.result_s3_object_id.is_not(None)
|
||
)
|
||
)
|
||
result_s3_ids = [row[0] for row in rows]
|
||
await session.execute(delete(RecognitionJob).where(RecognitionJob.id.in_(self.job_ids)))
|
||
if self._file_ids:
|
||
await session.execute(delete(StoredFile).where(StoredFile.id.in_(self._file_ids)))
|
||
s3_ids = [*self._s3_ids, *self._extra_s3_ids, *result_s3_ids]
|
||
if s3_ids:
|
||
await session.execute(delete(S3Object).where(S3Object.id.in_(s3_ids)))
|
||
if self._entity_ids:
|
||
await session.execute(delete(LegalEntity).where(LegalEntity.id.in_(self._entity_ids)))
|
||
self.job_ids.clear()
|
||
self._file_ids.clear()
|
||
self._s3_ids.clear()
|
||
self._extra_s3_ids.clear()
|
||
self._entity_ids.clear()
|