arm-core/api/tests/test_repositories_consolidations.py

"""Репозиторий консолидаций (Срез 4, задача 1.2): выборки версий, блокировки, строки свода, состояние файлов ЮЛ."""

import asyncio
import datetime
import uuid
from decimal import Decimal

import pytest
from sqlalchemy import select, text, update
from sqlalchemy.exc import DBAPIError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from src.domain.checklist import EntityFilesState
from src.domain.policy import VisibilityScope
from src.models.access import AppUser
from src.models.consolidation import ConsolidationDocument, ConsolidationVersion
from src.models.directories import CompanyGroup, LegalEntity, RegionalOffice
from src.models.enums import (
    ConsolidationApproach,
    ConsolidationStatus,
    DocumentStatus,
    DocumentType,
    FileStatus,
    OsvDataType,
    OsvGroupType,
    PeriodQuarter,
    RecognitionStatus,
)
from src.repositories import consolidations as repo
from tests.factories import (
    make_company_group,
    make_confirmed_osv,
    make_consolidation,
    make_document,
    make_legal_entity,
    make_lock,
    make_regional_office,
    make_stored_file,
    make_user,
)

GLOBAL = VisibilityScope(office_id=None)
IN_WORK = (
    ConsolidationStatus.INITIATED,
    ConsolidationStatus.DOCUMENTS_PROCESSING,
    ConsolidationStatus.DATA_CORRECTION,
)
LOCK_NOT_AVAILABLE_TIMEOUT = 5


async def _world(
    session: AsyncSession, *, members: int = 2
) -> tuple[RegionalOffice, CompanyGroup, list[LegalEntity], AppUser]:
    """Филиал, ГК с `members` ЮЛ (имена «ЮЛ-0», «ЮЛ-1», …) и пользователь-инициатор."""
    office = await make_regional_office(session)
    group = await make_company_group(session, office=office)
    entities = [await make_legal_entity(session, name=f"ЮЛ-{i}", group=group) for i in range(members)]
    user = await make_user(session, email=f"u{uuid.uuid4().hex[:8]}@armka.local", office=office)
    return office, group, entities, user


@pytest.mark.integration
async def test_list_periods_sorts_periods_and_versions_descending(db_session: AsyncSession) -> None:
    _, group, entities, user = await _world(db_session)
    q1, first = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q1, members=entities, initiator=user
    )
    q3, _ = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )
    second = ConsolidationVersion(
        consolidation_id=q1.id,
        version_no=2,
        status=ConsolidationStatus.INITIATED,
        scenario=first.scenario,
        elimination_rules_version=first.elimination_rules_version,
        initiator_id=user.id,
        initiator_office_id=user.regional_office_id,
    )
    db_session.add(second)
    await db_session.flush()

    periods = await repo.list_periods(db_session, GLOBAL, group.id)

    assert [(c.year, c.quarter) for c, _ in periods] == [(2024, PeriodQuarter.Q3), (2024, PeriodQuarter.Q1)]
    assert [[v.version_no for v in versions] for _, versions in periods] == [[1], [2, 1]]
    assert periods[0][0].id == q3.id


@pytest.mark.integration
async def test_list_periods_hides_foreign_office_group(db_session: AsyncSession) -> None:
    office, group, entities, user = await _world(db_session)
    other = await make_regional_office(db_session)
    await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q1, members=entities, initiator=user
    )

    assert len(await repo.list_periods(db_session, VisibilityScope(office_id=office.id), group.id)) == 1
    assert await repo.list_periods(db_session, VisibilityScope(office_id=other.id), group.id) == []


@pytest.mark.integration
async def test_get_version_respects_visibility(db_session: AsyncSession) -> None:
    office, group, entities, user = await _world(db_session)
    other = await make_regional_office(db_session)
    consolidation, version = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )

    found = await repo.get_version(db_session, VisibilityScope(office_id=office.id), version.id)
    assert found is not None
    assert (found[0].id, found[1].id) == (version.id, consolidation.id)
    assert await repo.get_version(db_session, GLOBAL, version.id) is not None
    assert await repo.get_version(db_session, VisibilityScope(office_id=other.id), version.id) is None  # чужой филиал
    assert await repo.get_version(db_session, GLOBAL, uuid.uuid4()) is None


@pytest.mark.integration
async def test_get_consolidation_for_group_by_period(db_session: AsyncSession) -> None:
    office, group, entities, user = await _world(db_session)
    other = await make_regional_office(db_session)
    consolidation, _ = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )
    own = VisibilityScope(office_id=office.id)

    found = await repo.get_consolidation_for_group(db_session, own, group.id, 2024, PeriodQuarter.Q3)
    assert found is not None
    assert found.id == consolidation.id
    assert await repo.get_consolidation_for_group(db_session, own, group.id, 2024, PeriodQuarter.Q2) is None
    assert (
        await repo.get_consolidation_for_group(
            db_session, VisibilityScope(office_id=other.id), group.id, 2024, PeriodQuarter.Q3
        )
        is None
    )


@pytest.mark.integration
async def test_lock_version_returns_version_or_none(db_session: AsyncSession) -> None:
    _, group, entities, user = await _world(db_session)
    _, version = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )
    # статус меняется в БД в обход ORM: объект в сессии остаётся прежним, `populate_existing` перечитывает строку
    await db_session.execute(
        update(ConsolidationVersion)
        .where(ConsolidationVersion.id == version.id)
        .values(status=ConsolidationStatus.CANCELLED)
        .execution_options(synchronize_session=False)
    )
    assert version.status == ConsolidationStatus.INITIATED

    locked = await repo.lock_version(db_session, version.id)

    assert locked is not None
    assert locked is version
    assert locked.status == ConsolidationStatus.CANCELLED
    assert await repo.lock_version(db_session, uuid.uuid4()) is None


@pytest.mark.integration
async def test_lock_versions_for_entity_returns_only_in_work_member_versions_by_id(db_session: AsyncSession) -> None:
    _, group, entities, user = await _world(db_session, members=3)
    first, second, outsider = entities
    created: dict[ConsolidationStatus, ConsolidationVersion] = {}
    for quarter, status in zip(
        (PeriodQuarter.Q1, PeriodQuarter.Q2, PeriodQuarter.Q3, PeriodQuarter.Q4),
        (
            ConsolidationStatus.INITIATED,
            ConsolidationStatus.DATA_CORRECTION,
            ConsolidationStatus.VERIFIED,
            ConsolidationStatus.DOCUMENTS_PROCESSING,
        ),
        strict=True,
    ):
        _, created[status] = await make_consolidation(
            db_session, group=group, year=2024, quarter=quarter, members=[first, second], initiator=user, status=status
        )
    # ещё одна версия «в работе», в периметре которой ЮЛ нет
    await make_consolidation(
        db_session, group=group, year=2023, quarter=PeriodQuarter.Q1, members=[second, outsider], initiator=user
    )

    versions = await repo.lock_versions_for_entity(db_session, first.id, IN_WORK)

    expected = sorted(
        created[status].id
        for status in (
            ConsolidationStatus.INITIATED,
            ConsolidationStatus.DATA_CORRECTION,
            ConsolidationStatus.DOCUMENTS_PROCESSING,
        )
    )
    assert [version.id for version in versions] == expected


@pytest.mark.integration
async def test_lock_versions_for_entity_locks_superset_but_returns_members_only(
    committed_sessionmaker: async_sessionmaker[AsyncSession],
) -> None:
    """R-2: версия ГК «в работе», где ЮЛ не член, блокируется (правка периметра ждёт), но не возвращается."""
    async with committed_sessionmaker() as setup, setup.begin():
        _, group, entities, user = await _world(setup, members=3)
        member_entity, other_entity, third = entities
        _, member_version = await make_consolidation(
            setup,
            group=group,
            year=2024,
            quarter=PeriodQuarter.Q1,
            members=[member_entity, other_entity],
            initiator=user,
        )
        _, superset_version = await make_consolidation(
            setup, group=group, year=2024, quarter=PeriodQuarter.Q2, members=[other_entity, third], initiator=user
        )
        foreign_group = await make_company_group(setup)
        _, foreign_version = await make_consolidation(
            setup,
            group=foreign_group,
            year=2024,
            quarter=PeriodQuarter.Q1,
            members=[await make_legal_entity(setup, group=foreign_group)] * 1,
            initiator=user,
        )
        ids = (member_entity.id, member_version.id, superset_version.id, foreign_version.id)
    entity_id, member_id, superset_id, foreign_id = ids

    async def try_lock(session: AsyncSession, version_id: uuid.UUID) -> bool:
        """`True` — строка заблокирована другой транзакцией (NOWAIT не получил блокировку)."""
        try:
            async with session.begin_nested():
                await session.execute(
                    text("SELECT id FROM consolidation_version WHERE id = :id FOR UPDATE NOWAIT"), {"id": version_id}
                )
        except DBAPIError:
            return True
        return False

    async with committed_sessionmaker() as locker, locker.begin(), committed_sessionmaker() as probe, probe.begin():
        returned = await asyncio.wait_for(
            repo.lock_versions_for_entity(locker, entity_id, IN_WORK), timeout=LOCK_NOT_AVAILABLE_TIMEOUT
        )
        assert [version.id for version in returned] == [member_id]
        assert await try_lock(probe, member_id) is True
        assert await try_lock(probe, superset_id) is True  # надмножество: ЮЛ не член, но версия той же ГК
        assert await try_lock(probe, foreign_id) is False  # версия чужой ГК не блокируется


@pytest.mark.integration
async def test_active_lock_ignores_released(db_session: AsyncSession) -> None:
    _, group, entities, user = await _world(db_session)
    _, version = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )
    assert await repo.active_lock(db_session, version.id) is None
    await make_lock(db_session, version=version, user=user, released=True)
    assert await repo.active_lock(db_session, version.id) is None
    active = await make_lock(db_session, version=version, user=user)

    found = await repo.active_lock(db_session, version.id)

    assert found is not None
    assert found.id == active.id


@pytest.mark.integration
async def test_members_sorted_by_entity_name(db_session: AsyncSession) -> None:
    office = await make_regional_office(db_session)
    group = await make_company_group(db_session, office=office)
    user = await make_user(db_session, email="m@armka.local", office=office)
    entities = [await make_legal_entity(db_session, name=name, group=group) for name in ("Яблоко", "Арбуз", "Мандарин")]
    _, version = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )

    rows = await repo.members(db_session, version.id)

    assert [entity.name for _, entity in rows] == ["Арбуз", "Мандарин", "Яблоко"]
    assert all(member.consolidation_version_id == version.id for member, _ in rows)


@pytest.mark.integration
async def test_history_in_chronological_order(db_session: AsyncSession) -> None:
    _, group, entities, user = await _world(db_session)
    _, version = await make_consolidation(
        db_session,
        group=group,
        year=2024,
        quarter=PeriodQuarter.Q3,
        members=entities,
        initiator=user,
        status=ConsolidationStatus.DATA_CORRECTION,
    )

    history = await repo.history(db_session, version.id)

    assert [(row.from_status, row.to_status) for row in history] == [
        (None, ConsolidationStatus.INITIATED),
        (ConsolidationStatus.INITIATED, ConsolidationStatus.DATA_CORRECTION),
    ]


async def _osv_rows() -> list[dict[str, object]]:
    return [
        {"name": "Контрагент А", "inn": "7701000001", "opening_debit": Decimal("100.00")},
        {"name": "Контрагент Б", "inn": None, "closing_credit": Decimal("5.50"), "data_type": OsvDataType.SUM},
        {"name": "Итого", "group_type": OsvGroupType.ACCOUNT},
        {"name": "Контрагент Н", "data_type": OsvDataType.NU},
        {"name": "Договор", "group_type": OsvGroupType.CONTRACT},
    ]


@pytest.mark.integration
async def test_load_osv_rows_keeps_counterparty_bu_sum_only_in_order(db_session: AsyncSession) -> None:
    office = await make_regional_office(db_session)
    group = await make_company_group(db_session, office=office)
    user = await make_user(db_session, email="o@armka.local", office=office)
    late = await make_legal_entity(db_session, name="Ящик", group=group)
    early = await make_legal_entity(db_session, name="Арбуз", group=group)
    _, version = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=[late, early], initiator=user
    )
    rows = await _osv_rows()
    docs = [
        await make_confirmed_osv(
            db_session, entity=entity, account=account, year=2024, quarter=quarter, is_cumulative=False, rows=rows
        )
        for entity, account, quarter in (
            (late, "60", PeriodQuarter.Q1),
            (early, "62", PeriodQuarter.Q1),
            (early, "60", PeriodQuarter.Q3),
            (early, "60", PeriodQuarter.Q1),
        )
    ]
    await make_confirmed_osv(  # не привязан к версии: в свод не попадает
        db_session, entity=early, account="60", year=2024, quarter=PeriodQuarter.Q2, is_cumulative=False, rows=rows
    )
    for document in docs:
        db_session.add(ConsolidationDocument(consolidation_version_id=version.id, document_id=document.id))
    await db_session.flush()

    loaded = await repo.load_osv_rows(db_session, version.id)

    # отброшены ACCOUNT, CONTRACT, NU; остались «А» (BU) и «Б» (SUM) каждого документа
    assert {entry.counterparty_name for entry, _, _ in loaded} == {"Контрагент А", "Контрагент Б"}
    assert len(loaded) == 2 * len(docs)
    keys = [(entity.name, doc.osv_account_code, doc.period_quarter, entry.row_number) for entry, doc, entity in loaded]
    assert keys == sorted(keys, key=lambda k: (k[0], k[1], k[2].value, k[3]))
    assert keys[0] == ("Арбуз", "60", PeriodQuarter.Q1, 1)
    assert [k[0] for k in keys].count("Ящик") == 2  # noqa: PLR2004


@pytest.mark.integration
async def test_bound_documents_and_document_versions(db_session: AsyncSession) -> None:
    _, group, entities, user = await _world(db_session)
    _, first = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )
    _, second = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q2, members=entities, initiator=user
    )
    document = await make_confirmed_osv(
        db_session, entity=entities[0], account="60", year=2024, quarter=PeriodQuarter.Q1, is_cumulative=False, rows=[]
    )
    unbound = await make_confirmed_osv(
        db_session, entity=entities[1], account="60", year=2024, quarter=PeriodQuarter.Q1, is_cumulative=False, rows=[]
    )
    for version in (first, second):
        db_session.add(ConsolidationDocument(consolidation_version_id=version.id, document_id=document.id))
    await db_session.flush()

    assert [d.id for d in await repo.bound_documents(db_session, first.id)] == [document.id]
    assert sorted(await repo.document_versions(db_session, document.id)) == sorted([first.id, second.id])
    assert await repo.document_versions(db_session, unbound.id) == []


@pytest.mark.integration
async def test_entity_files_state_in_progress_and_last_status(db_session: AsyncSession) -> None:
    group = await make_company_group(db_session)
    busy = await make_legal_entity(db_session, group=group)
    failed = await make_legal_entity(db_session, group=group)
    idle = await make_legal_entity(db_session, group=group)
    recovered = await make_legal_entity(db_session, group=group)
    await make_stored_file(db_session, entity=busy, status=FileStatus.UPLOADED)
    await make_stored_file(db_session, entity=busy, status=FileStatus.PROCESSING)
    old, _ = await make_stored_file(db_session, entity=failed, status=FileStatus.PROCESSED)
    newer, _ = await make_stored_file(db_session, entity=failed, status=FileStatus.ERROR)
    old.uploaded_at = datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC)
    newer.uploaded_at = datetime.datetime(2024, 2, 1, tzinfo=datetime.UTC)
    bad, _ = await make_stored_file(db_session, entity=recovered, status=FileStatus.ERROR)
    good, _ = await make_stored_file(db_session, entity=recovered, status=FileStatus.PROCESSED)
    bad.uploaded_at = datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC)
    good.uploaded_at = datetime.datetime(2024, 2, 1, tzinfo=datetime.UTC)
    await db_session.flush()

    state = await repo.entity_files_state(db_session, [busy.id, failed.id, idle.id, recovered.id])

    assert state[busy.id].has_in_progress is True
    assert state[failed.id] == EntityFilesState(has_in_progress=False, last_file_status=FileStatus.ERROR)
    assert state[idle.id] == EntityFilesState(has_in_progress=False, last_file_status=None)
    assert state[recovered.id] == EntityFilesState(has_in_progress=False, last_file_status=FileStatus.PROCESSED)
    assert await repo.entity_files_state(db_session, []) == {}


@pytest.mark.integration
async def test_rules_for_basic_and_extended_from_seed(seeded: AsyncSession) -> None:
    basic = await repo.rules_for(seeded, ConsolidationApproach.BASIC)
    extended = await repo.rules_for(seeded, ConsolidationApproach.EXTENDED)

    assert {(rule.document_type, account) for rule, account in basic} == {
        (DocumentType.BALANCE, None),
        (DocumentType.INCOME, None),
        *[(DocumentType.OSV, code) for code in ("01", "58", "60", "62", "66", "67", "76", "90")],
    }
    assert all(rule.approach == ConsolidationApproach.BASIC and rule.is_mandatory for rule, _ in basic)
    assert extended
    assert {rule.approach for rule, _ in extended} == {ConsolidationApproach.EXTENDED}


@pytest.mark.integration
async def test_unconfirmed_documents_filters_status_year_entity_and_superseded(db_session: AsyncSession) -> None:
    group = await make_company_group(db_session)
    entity = await make_legal_entity(db_session, group=group)
    other = await make_legal_entity(db_session, group=group)

    async def doc(
        owner: LegalEntity,
        *,
        status: DocumentStatus,
        year: int = 2024,
        file_status: FileStatus = FileStatus.PROCESSED,
        account: str = "60",
    ) -> uuid.UUID:
        file, job = await make_stored_file(
            db_session, entity=owner, status=file_status, job_status=RecognitionStatus.done
        )
        document = await make_document(
            db_session, file=file, job=job, account=account, year=year, quarter=PeriodQuarter.Q3, status=status
        )
        return document.id

    needs_check = await doc(entity, status=DocumentStatus.NEEDS_CHECK)
    recognized = await doc(entity, status=DocumentStatus.RECOGNIZED, account="58")
    await doc(entity, status=DocumentStatus.CONFIRMED)  # подтверждённый — не кандидат
    await doc(entity, status=DocumentStatus.ERROR)
    await doc(entity, status=DocumentStatus.NEEDS_CHECK, year=2023)
    await doc(entity, status=DocumentStatus.NEEDS_CHECK, file_status=FileStatus.SUPERSEDED)
    await doc(other, status=DocumentStatus.NEEDS_CHECK)

    found = await repo.unconfirmed_documents(db_session, [entity.id], 2024)

    assert sorted(d.id for d in found) == sorted([needs_check, recognized])
    assert await repo.unconfirmed_documents(db_session, [], 2024) == []


@pytest.mark.integration
async def test_repository_returns_selectable_models(db_session: AsyncSession) -> None:
    """Дымовая проверка: список версий ЮЛ после `make_consolidation` читается обычным запросом."""
    _, group, entities, user = await _world(db_session)
    _, version = await make_consolidation(
        db_session, group=group, year=2024, quarter=PeriodQuarter.Q3, members=entities, initiator=user
    )
    assert (await db_session.execute(select(ConsolidationVersion.id))).scalars().all().count(version.id) == 1