arm-core/api/tests/factories.py
Andrey Epifancev 3b4d3a45fe fix(exchange): INVALID_REQUEST публикатора переводит файл в ERROR
T4.1: публикатор, переводя задание в финальный error/INVALID_REQUEST, теперь в той же транзакции (под тем же
FOR UPDATE) ставит файлу текущего задания ERROR (spec primary-files «Жизненный цикл статуса файла»); файл
чужого задания не меняется. Тест test_invalid_request_marks_job_error_and_pass_continues[current|non_current];
CommittedJobs.make_file_job принимает original_name.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012vJkmBGKKqMoEMzvsXG3y2
2026-10-05 22:13:58 +04:00

548 lines
23 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Фабрики тестовых данных для БД-тестов.
- `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`.
- `make_stored_file`, `make_document`, `make_result_object` — цепочка файла/задания/документа и объект результата в
`FakeStorage` для тестов загрузки и маппинга (Срез 2c; `flush`, без `commit`). `make_entries_from_result` создаёт
Task 9 (R-2).
- `CommittedJobs` — задания распознавания, ЗАФИКСИРОВАННЫЕ в БД (`commit`): для фоновых процессов обмена, которые
работают в собственных транзакциях и не видят чужую незафиксированную транзакцию. `make_file_job` / `add_job` — файл
с текущим заданием и второе задание того же файла (статусы `stored_file` в обмене, Срез 2c). Всё созданное удаляется
`cleanup()` в порядке FK, включая строки `s3_object` файлов результата, зарегистрированные inbox'ом.
"""
import datetime
import decimal
import json
import secrets
import uuid
from collections.abc import Sequence
from typing import Any
from sqlalchemy import delete, func, select, update
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,
PeriodQuarter,
RecognitionPipeline,
RecognitionStatus,
RelatedGroupKind,
S3BucketRole,
)
from src.models.files import Document, RecognitionJob, S3Object, StoredFile
from src.models.osv import OsvEntry
from tests.conftest import FakeStorage
DEFAULT_FILE_TYPES = {RecognitionPipeline.tabular: FileType.XLSX, RecognitionPipeline.image: FileType.PDF}
TABULAR_FILE_TYPES = frozenset({FileType.XLS, FileType.XLSX, FileType.XLSB})
FACTORY_BUCKET = "arm-ka"
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_stored_file( # noqa: PLR0913 - сигнатура фабрики фиксирована планом
session: AsyncSession,
*,
entity: LegalEntity,
status: FileStatus = FileStatus.UPLOADED,
file_type: FileType = FileType.XLS,
original_name: str = "osv.xls",
content_hash: str | None = None,
job_status: RecognitionStatus = RecognitionStatus.queued,
s3_key: str | None = None,
uploaded_by: uuid.UUID | None = None,
) -> tuple[StoredFile, RecognitionJob]:
"""`s3_object` (uploads) + `stored_file` + `recognition_job` (attempt 1), `current_recognition_job_id` = задание.
Конвейер задания выводится из типа файла (XLS/XLSX/XLSB — tabular, иначе image). Данные остаются в транзакции
сессии (`flush`).
"""
now = _now()
s3 = S3Object(
bucket_role=S3BucketRole.uploads,
bucket=FACTORY_BUCKET,
key=s3_key or f"uploads/{uuid.uuid4()}/{original_name}",
created_at=now,
)
session.add(s3)
await session.flush()
file = StoredFile(
legal_entity_id=entity.id,
file_type=file_type,
original_name=original_name,
s3_object_id=s3.id,
content_hash=content_hash,
status=status,
uploaded_by=uploaded_by,
uploaded_at=now,
)
session.add(file)
await session.flush()
pipeline = RecognitionPipeline.tabular if file_type in TABULAR_FILE_TYPES else RecognitionPipeline.image
job = RecognitionJob(
file_id=file.id,
pipeline=pipeline,
request_subject=f"recognition.{pipeline.value}.request",
schema_version="1.0",
status=job_status,
)
session.add(job)
await session.flush()
file.current_recognition_job_id = job.id
await session.flush()
return file, job
async def make_document( # noqa: PLR0913 - сигнатура фабрики фиксирована планом
session: AsyncSession,
*,
file: StoredFile,
job: RecognitionJob,
account: str = "60",
year: int = 2025,
quarter: PeriodQuarter = PeriodQuarter.Q1,
status: DocumentStatus = DocumentStatus.RECOGNIZED,
is_cumulative: bool = True,
) -> Document:
"""ОСВ-документ по счёту `account` файла `file` (`flush`).
`job` — задание, породившее документ (у `document` FK на задание нет; параметр держит вызовы единообразными с
`make_stored_file`).
"""
assert job.file_id == file.id # noqa: S101 - защита от перепутанных аргументов в тестах
now = _now()
document = Document(
stored_file_id=file.id,
legal_entity_id=file.legal_entity_id,
document_type=DocumentType.OSV,
osv_account_code=account,
period_year=year,
period_quarter=quarter,
is_cumulative=is_cumulative,
status=status,
created_at=now,
updated_at=now,
)
session.add(document)
await session.flush()
return document
async def make_result_object(
session: AsyncSession, *, job: RecognitionJob, storage: FakeStorage, result: dict[str, Any]
) -> S3Object:
"""JSON результата — в `storage` под `results/{file_id}/{job_id}/result.json` + `s3_object` + ссылка в задании."""
key = f"results/{job.file_id}/{job.id}/result.json"
payload = json.dumps(result, ensure_ascii=False).encode("utf-8")
await storage.put_object(key, payload)
s3 = S3Object(
bucket_role=S3BucketRole.results,
bucket=storage.bucket,
key=key,
content_type="application/json",
size_bytes=len(payload),
created_at=_now(),
)
session.add(s3)
await session.flush()
job.result_s3_object_id = s3.id
await session.flush()
return s3
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 make_file_job(
self,
*,
file_status: FileStatus = FileStatus.UPLOADED,
job_status: RecognitionStatus = RecognitionStatus.queued,
original_name: str | None = "osv.xls",
**job_fields: Any,
) -> tuple[uuid.UUID, uuid.UUID]:
"""Файл с текущим заданием (`make_stored_file`: `current_recognition_job_id` = задание), зафиксированный в БД.
`job_fields` — поля `recognition_job` поверх фабрики (`reply_deadline`, `published_at`, `metrics`, …);
`original_name=None` — файл без имени (запрос не собирается по контракту). Возвращает
`(stored_file.id, recognition_job.id)`.
"""
async with self._session_factory() as session, session.begin():
entity = await make_legal_entity(session, name="ООО Тест статусов файла")
file, job = await make_stored_file(
session,
entity=entity,
status=file_status,
job_status=job_status,
original_name=original_name or "file",
)
file.original_name = original_name
for name, value in job_fields.items():
setattr(job, name, value)
await session.flush()
self._entity_ids.append(entity.id)
self._s3_ids.append(file.s3_object_id)
self._file_ids.append(file.id)
self.job_ids.append(job.id)
return file.id, job.id
async def add_job(
self,
file_id: uuid.UUID,
*,
status: RecognitionStatus = RecognitionStatus.queued,
current: bool = True,
**job_fields: Any,
) -> uuid.UUID:
"""Ещё одно задание файла `file_id` (как replay: `attempt` = максимум по файлу + 1); `current=True` — оно
становится текущим заданием файла."""
async with self._session_factory() as session, session.begin():
file = await session.get(StoredFile, file_id)
assert file is not None
last = (
await session.execute(select(func.max(RecognitionJob.attempt)).where(RecognitionJob.file_id == file_id))
).scalar_one()
job_fields.setdefault("attempt", (last or 0) + 1)
tabular = file.file_type in TABULAR_FILE_TYPES
pipeline = RecognitionPipeline.tabular if tabular else RecognitionPipeline.image
job = RecognitionJob(
file_id=file_id,
pipeline=pipeline,
request_subject=f"recognition.{pipeline.value}.request",
schema_version="1.0",
status=status,
**job_fields,
)
session.add(job)
await session.flush()
if current:
file.current_recognition_job_id = job.id
self.job_ids.append(job.id)
return job.id
async def file(self, file_id: uuid.UUID) -> StoredFile:
"""Свежее состояние файла из БД (новая сессия)."""
async with self._session_factory() as session:
file = await session.get(StoredFile, file_id)
assert file is not None
return file
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: строки ОСВ и документы файлов → задания → файлы → объекты S3 и ЮЛ."""
async with self._session_factory() as session, session.begin():
result_s3_ids: list[uuid.UUID] = []
if self._file_ids:
# задания ссылаются на файл, файл — на текущее задание (use_alter): сначала разрываем обратную ссылку
await session.execute(
update(StoredFile).where(StoredFile.id.in_(self._file_ids)).values(current_recognition_job_id=None)
)
documents = select(Document.id).where(Document.stored_file_id.in_(self._file_ids))
await session.execute(delete(OsvEntry).where(OsvEntry.document_id.in_(documents)))
await session.execute(delete(Document).where(Document.stored_file_id.in_(self._file_ids)))
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(OsvEntry).where(OsvEntry.recognition_job_id.in_(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)))
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()