"""Фабрики тестовых данных для БД-тестов. - `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()