arm-core/api/tests/factories.py

"""Фабрики тестовых данных для БД-тестов.

- `make_osv_entry` — минимальная цепочка FK поверх сессии (`flush`, без `commit`): откатывается фикстурой `db_session`.
- `CommittedJobs` — задания распознавания, ЗАФИКСИРОВАННЫЕ в БД (`commit`): для фоновых процессов обмена, которые
  работают в собственных транзакциях и не видят чужую незафиксированную транзакцию. Всё созданное удаляется
  `cleanup()` в порядке FK.
"""

import datetime
import decimal
import uuid
from typing import Any

from sqlalchemy import delete
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._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,
    ) -> 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,
            )
            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

    async def cleanup(self) -> None:
        """Удаляет всё созданное в порядке FK."""
        async with self._session_factory() as session, session.begin():
            if self.job_ids:
                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)))
            if self._s3_ids:
                await session.execute(delete(S3Object).where(S3Object.id.in_(self._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._entity_ids.clear()