"""Фабрики тестовых данных для БД-тестов. - `make_regional_office`, `make_company_group`, `make_legal_entity`, `make_related_group`, `make_membership`, `make_user` — справочники Среза 1 (`flush`, без `commit`): откатываются фикстурой `db_session`. Идентификаторы по умолчанию случайные (ИНН, коды, имена филиалов), поэтому фабрики можно вызывать несколько раз в одном тесте. - `make_osv_entry` — минимальная цепочка FK поверх сессии (`flush`, без `commit`): откатывается фикстурой `db_session`. - `CommittedJobs` — задания распознавания, ЗАФИКСИРОВАННЫЕ в БД (`commit`): для фоновых процессов обмена, которые работают в собственных транзакциях и не видят чужую незафиксированную транзакцию. Всё созданное удаляется `cleanup()` в порядке FK, включая строки `s3_object` файлов результата, зарегистрированные inbox'ом. """ import datetime import decimal import secrets import uuid from collections.abc import Sequence from typing import Any from sqlalchemy import delete, func, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from src.models.access import AppUser, Role, UserRole from src.models.directories import ( CompanyGroup, LegalEntity, RegionalOffice, RelatedGroup, RelatedGroupMember, ) from src.models.enums import ( DocumentStatus, DocumentType, EntityType, FileStatus, FileType, RecognitionPipeline, RecognitionStatus, RelatedGroupKind, 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} def _now() -> datetime.datetime: return datetime.datetime.now(datetime.UTC) async def make_regional_office( session: AsyncSession, *, name: str | None = None, is_head_office: bool = False ) -> RegionalOffice: """Филиал; имя по умолчанию уникальное (`regional_office.name` UNIQUE).""" office = RegionalOffice(name=name or f"Филиал {uuid.uuid4().hex[:8]}", is_head_office=is_head_office) session.add(office) await session.flush() return office async def make_company_group( session: AsyncSession, *, name: str = "ГК Тест", office: RegionalOffice | None = None, crm_id: str | None = None, code: str | None = None, ) -> CompanyGroup: """ГК; `code` по умолчанию `GK-T<6 hex>`, `valid_from` — сейчас.""" group = CompanyGroup( code=code or f"GK-T{uuid.uuid4().hex[:6]}", crm_id=crm_id, name=name, base_regional_office_id=office.id if office is not None else None, valid_from=_now(), ) session.add(group) await session.flush() return group async def make_legal_entity( # noqa: PLR0913 - сигнатура фабрики фиксирована планом session: AsyncSession, *, name: str = "ООО Тест", inn: str | None = None, entity_type: EntityType = EntityType.ЮЛ, group: CompanyGroup | None = None, office: RegionalOffice | None = None, close_date: datetime.date | None = None, ) -> LegalEntity: """ЮЛ; ИНН по умолчанию — 10 случайных цифр, филиал по умолчанию — базовый филиал группы.""" office_id = office.id if office is not None else (group.base_regional_office_id if group is not None else None) entity = LegalEntity( inn=inn or "".join(secrets.choice("0123456789") for _ in range(10)), entity_type=entity_type, name=name, company_group_id=group.id if group is not None else None, regional_office_id=office_id, close_date=close_date, valid_from=_now(), ) session.add(entity) await session.flush() return entity async def make_related_group( session: AsyncSession, *, group: CompanyGroup, kind: RelatedGroupKind = RelatedGroupKind.gsk, name: str = "ГСК Тест" ) -> RelatedGroup: """ГСК/ГСЗ группы `group`; код `GSK-T…`/`GSZ-T…`.""" related = RelatedGroup( code=f"{kind.value.upper()}-T{uuid.uuid4().hex[:6]}", company_group_id=group.id, kind=kind, name=name, created_date=_now().date(), ) session.add(related) await session.flush() return related async def make_membership( session: AsyncSession, *, related_group: RelatedGroup, entity: LegalEntity, is_current: bool = True ) -> RelatedGroupMember: """Членство ЮЛ в ГСК/ГСЗ; `kind` копируется из группы (v5.14).""" now = _now() member = RelatedGroupMember( related_group_id=related_group.id, legal_entity_id=entity.id, kind=related_group.kind, valid_from=now, valid_to=None if is_current else now, is_current=is_current, ) session.add(member) await session.flush() return member async def make_user( # noqa: PLR0913 - сигнатура фабрики фиксирована планом session: AsyncSession, *, email: str, office: RegionalOffice, roles: Sequence[str] = (), full_name: str = "Тест Тестов", is_active: bool = True, ) -> AppUser: """Пользователь с ролями по `Role.code` (роли должны быть засеяны: фикстура `seeded`).""" user = AppUser(email=email, full_name=full_name, regional_office_id=office.id, is_active=is_active) session.add(user) await session.flush() for code in roles: role_id = (await session.execute(select(Role.id).where(Role.code == code))).scalar_one() session.add(UserRole(user_id=user.id, role_id=role_id)) await session.flush() return user 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 async def metrics_sql_state(self, job_id: uuid.UUID) -> tuple[bool, str | None]: """Как хранится `recognition_job.metrics` на уровне SQL: (`IS NULL`, `jsonb_typeof`). SQL NULL — `(True, None)`; jsonb `null` — `(False, "null")`; объект — `(False, "object")`. """ async with self._session_factory() as session: row = ( await session.execute( select(RecognitionJob.metrics.is_(None), func.jsonb_typeof(RecognitionJob.metrics)).where( RecognitionJob.id == job_id ) ) ).one() return bool(row[0]), row[1] 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()