b4522addc0
- consolidation_version.elimination_rules_version NOT NULL: миграция 0010 проставляет существующим версиям номер правил, действовавших на дату создания (по Москве), иначе последних открытых; без правил — отказ; создание версии без действующих правил — 409 ELIMINATION_RULES_MISSING; ветки «правила не зафиксированы» убраны (rules, preset, is_version_fixed всегда true); - GET …/elimination-totals: все тройки «правило ББ × упорядоченная пара периметра» с ΣA, ΣB, расхождением, суммой элиминации, статусом AUTO/DISCREPANCY/ONE_SIDED/MANUAL и osv_entry_ids; сводная ОСВ строится один раз; сбор строк сторон вынесен в _collect_rule_sides — общий с упаковкой тройки. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
866 lines
38 KiB
Python
866 lines
38 KiB
Python
"""Фабрики тестовых данных для БД-тестов.
|
||
|
||
- `make_regional_office`, `make_company_group`, `make_legal_entity`, `make_related_group`, `make_membership`,
|
||
`make_user` — справочники Среза 1 (`flush`, без `commit`): откатываются фикстурой `db_session`. Идентификаторы по
|
||
умолчанию случайные (ИНН, коды, имена филиалов), поэтому фабрики можно вызывать несколько раз в одном тесте.
|
||
- `make_consolidation`, `make_version`, `make_lock`, `make_confirmed_osv` — консолидация ГК с версией №1, периметром и
|
||
историей; ещё одна версия; блокировка версии; подтверждённая ОСВ (файл `PROCESSED` + задание `done` + документ
|
||
`CONFIRMED` + строки контрагентов) для тестов Среза 4 (`flush`, без `commit`).
|
||
- `make_osv_entry` — минимальная цепочка FK поверх сессии (`flush`, без `commit`): откатывается фикстурой `db_session`.
|
||
- `make_stored_file`, `make_document`, `make_result_object` — цепочка файла/задания/документа и объект результата в
|
||
`FakeStorage` для тестов загрузки и маппинга (Срез 2c; `flush`, без `commit`). `make_entries_from_result` —
|
||
документы и строки ОСВ файла из JSON результата (golden-снимок) теми же функциями, что и цикл маппинга (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, insert, select, update
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||
|
||
from src.contract import request_subject
|
||
from src.contract.result import OsvDocument
|
||
from src.contract.versions import parse_result
|
||
from src.domain.consolidation_periods import reporting_date
|
||
from src.domain.osv_mapping import build_entries, document_fields
|
||
from src.domain.policy import FINISHED_STATUSES
|
||
from src.models.access import AppUser, Role, UserRole
|
||
from src.models.accounting import ChartOfAccount
|
||
from src.models.consolidation import (
|
||
Consolidation,
|
||
ConsolidationDocument,
|
||
ConsolidationLock,
|
||
ConsolidationMember,
|
||
ConsolidationStatusHistory,
|
||
ConsolidationVersion,
|
||
)
|
||
from src.models.directories import (
|
||
CompanyGroup,
|
||
LegalEntity,
|
||
RegionalOffice,
|
||
RelatedGroup,
|
||
RelatedGroupMember,
|
||
)
|
||
from src.models.enums import (
|
||
ConsolidationApproach,
|
||
ConsolidationScenario,
|
||
ConsolidationStatus,
|
||
DocumentCheckResult,
|
||
DocumentIssueKind,
|
||
DocumentStatus,
|
||
DocumentType,
|
||
EntityType,
|
||
FileStatus,
|
||
FileType,
|
||
IssueSeverity,
|
||
IssueSource,
|
||
OsvDataType,
|
||
OsvGroupType,
|
||
PeriodQuarter,
|
||
RecognitionPipeline,
|
||
RecognitionStatus,
|
||
RelatedGroupKind,
|
||
S3BucketRole,
|
||
)
|
||
from src.models.files import Document, RecognitionJob, S3Object, StoredFile
|
||
from src.models.osv import OsvEntry
|
||
from src.models.statements import DocumentIssue
|
||
from src.services.result_mapping import MAPPING_WARNINGS_KEY, _entry_values, _group_schema
|
||
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()
|
||
file_id = uuid.uuid4()
|
||
s3 = S3Object(
|
||
bucket_role=S3BucketRole.uploads,
|
||
bucket=FACTORY_BUCKET,
|
||
key=s3_key or f"uploads/{file_id}/{uuid.uuid4()}/{original_name}", # как `upload_key`: папка файла
|
||
created_at=now,
|
||
)
|
||
session.add(s3)
|
||
await session.flush()
|
||
file = StoredFile(
|
||
id=file_id,
|
||
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=request_subject(pipeline.value),
|
||
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()
|
||
# v5.19: file_ordinal уникален в файле — берём следующий свободный
|
||
ordinal = (
|
||
await session.execute(
|
||
select(func.coalesce(func.max(Document.file_ordinal) + 1, 0)).where(Document.stored_file_id == file.id)
|
||
)
|
||
).scalar_one()
|
||
document = Document(
|
||
stored_file_id=file.id,
|
||
legal_entity_id=file.legal_entity_id,
|
||
document_type=DocumentType.OSV,
|
||
file_ordinal=ordinal,
|
||
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_document_issue(session: AsyncSession, *, document_id: uuid.UUID, code: str = "MAPPING_WARNING") -> Any:
|
||
"""Запись `document_issue` (предупреждение ядра уровня `info`) документа (`flush`)."""
|
||
issue = DocumentIssue(
|
||
document_id=document_id,
|
||
kind=DocumentIssueKind.warning,
|
||
code=code,
|
||
severity=IssueSeverity.info,
|
||
message="тестовое предупреждение",
|
||
source=IssueSource.core,
|
||
created_at=_now(),
|
||
)
|
||
session.add(issue)
|
||
await session.flush()
|
||
return issue
|
||
|
||
|
||
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_entries_from_result(
|
||
session: AsyncSession,
|
||
file: StoredFile,
|
||
job: RecognitionJob,
|
||
result_json: dict[str, Any],
|
||
*,
|
||
status: DocumentStatus = DocumentStatus.NEEDS_CHECK,
|
||
) -> list[Document]:
|
||
"""Документ на каждый лист ОСВ результата + его строки `osv_entry` (`flush`), `job.group_schema` и
|
||
`metrics.mapping_warnings = []` — как после цикла маппинга, но без S3, резолва контрагентов и автозамены.
|
||
|
||
Строки строит `build_entries` (без `TOTAL`, `parent_entry_id` по `parent_row_number`, `account_id` по плану счетов
|
||
сессии), поля строки — `_entry_values` маппера. У golden-снимка нет `job_id`/`file_id`/`service.version` —
|
||
подставляются. Статус документов — `status` (по умолчанию `NEEDS_CHECK`), проверка периода — `UNKNOWN`.
|
||
"""
|
||
assert job.file_id == file.id # noqa: S101 - защита от перепутанных аргументов в тестах
|
||
payload = {**result_json, "job_id": str(job.id), "file_id": str(file.id)}
|
||
payload["service"] = {"version": "test", **result_json.get("service", {"name": "tabular-parser"})}
|
||
result = parse_result(json.dumps(payload).encode()) # модель — по MAJOR `schema_version` (снимки 1.x и 2.0)
|
||
sheets = [doc for doc in result.documents if isinstance(doc, OsvDocument)]
|
||
chart = {
|
||
code: account_id for code, account_id in await session.execute(select(ChartOfAccount.code, ChartOfAccount.id))
|
||
}
|
||
now = _now()
|
||
documents: list[Document] = []
|
||
entries: list[dict[str, Any]] = []
|
||
for sheet in sheets:
|
||
fields = document_fields(sheet)
|
||
document = Document(
|
||
id=uuid.uuid4(),
|
||
stored_file_id=file.id,
|
||
legal_entity_id=file.legal_entity_id,
|
||
document_type=DocumentType.OSV,
|
||
file_ordinal=len(documents),
|
||
osv_account_code=fields.osv_account_code,
|
||
period_year=fields.period_year,
|
||
period_quarter=fields.period_quarter,
|
||
is_cumulative=fields.is_cumulative,
|
||
locator=fields.locator,
|
||
status=status,
|
||
period_check=DocumentCheckResult.UNKNOWN,
|
||
counterparty_check=DocumentCheckResult.UNKNOWN,
|
||
created_at=now,
|
||
updated_at=now,
|
||
)
|
||
documents.append(document)
|
||
entries.extend(
|
||
_entry_values(entry, document_id=document.id, job_id=job.id, owner_id=file.legal_entity_id)
|
||
for entry in build_entries(sheet, chart).entries
|
||
)
|
||
session.add_all(documents)
|
||
await session.flush()
|
||
if entries:
|
||
await session.execute(insert(OsvEntry).values(entries))
|
||
job.group_schema = [_group_schema(sheet) for sheet in sheets]
|
||
metrics = dict(job.metrics) if isinstance(job.metrics, dict) else {}
|
||
job.metrics = {**metrics, MAPPING_WARNINGS_KEY: metrics.get(MAPPING_WARNINGS_KEY, [])}
|
||
await session.flush()
|
||
return documents
|
||
|
||
|
||
async def make_version( # noqa: PLR0913 - сигнатура фабрики фиксирована планом
|
||
session: AsyncSession,
|
||
*,
|
||
consolidation: Consolidation,
|
||
version_no: int,
|
||
members: Sequence[LegalEntity],
|
||
initiator: AppUser,
|
||
status: ConsolidationStatus = ConsolidationStatus.INITIATED,
|
||
) -> ConsolidationVersion:
|
||
"""Версия `version_no` консолидации: `consolidation_member` на каждое ЮЛ и история (NULL → INITIATED[→ status]).
|
||
|
||
`is_in_work = false` для VERIFIED / CANCELLED / ARCHIVED / REJECTED. История: если `status` не `INITIATED`, второй
|
||
записью — переход `INITIATED → status` (одним пользователем). Данные остаются в транзакции сессии (`flush`).
|
||
"""
|
||
now = _now()
|
||
version = ConsolidationVersion(
|
||
consolidation_id=consolidation.id,
|
||
version_no=version_no,
|
||
status=status,
|
||
scenario=ConsolidationScenario.CREATE,
|
||
approach=ConsolidationApproach.BASIC,
|
||
goodwill_enabled=False,
|
||
elimination_rules_version=1, # seed правил §4 — версия v1; с DBML v5.23 поле обязательно
|
||
initiator_id=initiator.id,
|
||
initiator_office_id=initiator.regional_office_id,
|
||
is_in_work=status not in FINISHED_STATUSES,
|
||
created_at=now,
|
||
created_by=initiator.id,
|
||
)
|
||
session.add(version)
|
||
await session.flush()
|
||
session.add_all(
|
||
ConsolidationMember(consolidation_version_id=version.id, legal_entity_id=entity.id, is_head_entity=False)
|
||
for entity in members
|
||
)
|
||
history = [(None, ConsolidationStatus.INITIATED)]
|
||
if status != ConsolidationStatus.INITIATED:
|
||
history.append((ConsolidationStatus.INITIATED, status))
|
||
session.add_all(
|
||
ConsolidationStatusHistory(
|
||
version_id=version.id,
|
||
from_status=old,
|
||
to_status=new,
|
||
changed_by=initiator.id,
|
||
changed_at=now + datetime.timedelta(milliseconds=position), # порядок записей детерминирован
|
||
)
|
||
for position, (old, new) in enumerate(history)
|
||
)
|
||
await session.flush()
|
||
return version
|
||
|
||
|
||
async def make_consolidation( # noqa: PLR0913 - сигнатура фабрики фиксирована планом
|
||
session: AsyncSession,
|
||
*,
|
||
group: CompanyGroup,
|
||
year: int,
|
||
quarter: PeriodQuarter,
|
||
members: Sequence[LegalEntity],
|
||
initiator: AppUser,
|
||
status: ConsolidationStatus = ConsolidationStatus.INITIATED,
|
||
) -> tuple[Consolidation, ConsolidationVersion]:
|
||
"""Консолидация ГК за период и её версия №1 (периметр `members`, история первой записью); `flush`.
|
||
|
||
`reporting_date` — по правилу спеки: Q1 → 01.04, Q2 → 01.07, Q3 → 01.10, Q4 → 01.01 следующего года.
|
||
"""
|
||
consolidation = Consolidation(
|
||
company_group_id=group.id,
|
||
year=year,
|
||
quarter=quarter,
|
||
reporting_date=reporting_date(year, quarter),
|
||
created_at=_now(),
|
||
)
|
||
session.add(consolidation)
|
||
await session.flush()
|
||
version = await make_version(
|
||
session, consolidation=consolidation, version_no=1, members=members, initiator=initiator, status=status
|
||
)
|
||
return consolidation, version
|
||
|
||
|
||
async def make_lock(
|
||
session: AsyncSession, *, version: ConsolidationVersion, user: AppUser, released: bool = False
|
||
) -> ConsolidationLock:
|
||
"""Блокировка версии пользователем `user`; `released=True` — уже снята (`released_at` задан); `flush`."""
|
||
now = _now()
|
||
lock = ConsolidationLock(
|
||
consolidation_version_id=version.id,
|
||
locked_by=user.id,
|
||
locked_at=now,
|
||
released_at=now if released else None,
|
||
)
|
||
session.add(lock)
|
||
await session.flush()
|
||
return lock
|
||
|
||
|
||
async def make_confirmed_osv( # noqa: PLR0913 - сигнатура фабрики фиксирована планом
|
||
session: AsyncSession,
|
||
*,
|
||
entity: LegalEntity,
|
||
account: str,
|
||
year: int,
|
||
quarter: PeriodQuarter,
|
||
is_cumulative: bool,
|
||
rows: Sequence[dict[str, Any]],
|
||
confirmed_by: uuid.UUID | None = None,
|
||
) -> Document:
|
||
"""Подтверждённая ОСВ ЮЛ: файл `PROCESSED` + задание `done` + документ `CONFIRMED` + строки; `flush`.
|
||
|
||
Каждая строка `rows` — словарь: `name` → `counterparty_name`, `inn` → `counterparty_inn`, остальные ключи — поля
|
||
`OsvEntry` как есть (`opening_debit`, `closing_credit`, … и при необходимости `group_type` / `data_type`, по
|
||
умолчанию `COUNTERPARTY` / `BU`). `row_number` — по порядку с 1; `account_code` / `account_root` — из `account`.
|
||
"""
|
||
file, job = await make_stored_file(
|
||
session, entity=entity, status=FileStatus.PROCESSED, job_status=RecognitionStatus.done
|
||
)
|
||
document = await make_document(
|
||
session,
|
||
file=file,
|
||
job=job,
|
||
account=account,
|
||
year=year,
|
||
quarter=quarter,
|
||
status=DocumentStatus.CONFIRMED,
|
||
is_cumulative=is_cumulative,
|
||
)
|
||
document.confirmed_by = confirmed_by
|
||
document.confirmed_at = _now()
|
||
entries: list[OsvEntry] = []
|
||
for number, row in enumerate(rows, start=1):
|
||
fields = dict(row)
|
||
fields.setdefault("group_type", OsvGroupType.COUNTERPARTY)
|
||
fields.setdefault("data_type", OsvDataType.BU)
|
||
fields["counterparty_name"] = fields.pop("name", None)
|
||
fields["counterparty_inn"] = fields.pop("inn", None)
|
||
entries.append(
|
||
OsvEntry(
|
||
recognition_job_id=job.id,
|
||
document_id=document.id,
|
||
owner_legal_entity_id=entity.id,
|
||
row_number=number,
|
||
account_code=account,
|
||
account_root=account.split(".", maxsplit=1)[0],
|
||
**fields,
|
||
)
|
||
)
|
||
session.add_all(entries)
|
||
await session.flush()
|
||
return document
|
||
|
||
|
||
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="files.tabular_uploaded",
|
||
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=request_subject(pipeline.value),
|
||
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=request_subject(pipeline.value),
|
||
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 _delete_consolidations(self, session: AsyncSession) -> None:
|
||
"""Привязки, блокировки, история, периметры, версии и консолидации, ссылающиеся на созданные ЮЛ и документы.
|
||
|
||
Версии — те, где созданное ЮЛ в периметре или привязан документ созданного файла; консолидации — только те,
|
||
у которых после этого не осталось версий. Выполняется до удаления документов и ЮЛ (FK на них).
|
||
"""
|
||
documents = select(Document.id).where(Document.stored_file_id.in_(self._file_ids))
|
||
by_document = select(ConsolidationDocument.consolidation_version_id).where(
|
||
ConsolidationDocument.document_id.in_(documents)
|
||
)
|
||
by_member = select(ConsolidationMember.consolidation_version_id).where(
|
||
ConsolidationMember.legal_entity_id.in_(self._entity_ids)
|
||
)
|
||
version_ids = [
|
||
row[0]
|
||
for row in await session.execute(
|
||
select(ConsolidationVersion.id).where(ConsolidationVersion.id.in_(by_document.union(by_member)))
|
||
)
|
||
]
|
||
if not version_ids:
|
||
return
|
||
consolidation_ids = select(ConsolidationVersion.consolidation_id).where(
|
||
ConsolidationVersion.id.in_(version_ids)
|
||
)
|
||
touched = [row[0] for row in await session.execute(consolidation_ids)]
|
||
await session.execute(
|
||
update(ConsolidationVersion)
|
||
.where(ConsolidationVersion.previous_version_id.in_(version_ids))
|
||
.values(previous_version_id=None)
|
||
)
|
||
await session.execute(
|
||
delete(ConsolidationDocument).where(ConsolidationDocument.consolidation_version_id.in_(version_ids))
|
||
)
|
||
await session.execute(
|
||
delete(ConsolidationLock).where(ConsolidationLock.consolidation_version_id.in_(version_ids))
|
||
)
|
||
await session.execute(
|
||
delete(ConsolidationStatusHistory).where(ConsolidationStatusHistory.version_id.in_(version_ids))
|
||
)
|
||
await session.execute(
|
||
delete(ConsolidationMember).where(ConsolidationMember.consolidation_version_id.in_(version_ids))
|
||
)
|
||
await session.execute(delete(ConsolidationVersion).where(ConsolidationVersion.id.in_(version_ids)))
|
||
remaining = select(ConsolidationVersion.consolidation_id).where(
|
||
ConsolidationVersion.consolidation_id.in_(touched)
|
||
)
|
||
await session.execute(
|
||
delete(Consolidation).where(Consolidation.id.in_(touched), Consolidation.id.not_in(remaining))
|
||
)
|
||
|
||
async def cleanup(self) -> None:
|
||
"""Удаляет всё созданное в порядке FK: строки ОСВ и документы файлов → задания → файлы → объекты S3 и ЮЛ."""
|
||
async with self._session_factory() as session, session.begin():
|
||
result_s3_ids: list[uuid.UUID] = []
|
||
await self._delete_consolidations(session)
|
||
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()
|