arm-core/api/tests/test_exchange_lifespan_nats.py

"""Обмен `recognition_job` в реальном lifespan и два инстанса ядра (финальное ревью I-7; spec recognition-job-exchange,
«Соединение с NATS по шаблону, без JetStream», «Inbox принимает результат идемпотентно»; design D-6).

Интеграционные тесты против NATS и PostgreSQL из `.env.test` (БД сессии pytest; задания фиксируются `CommittedJobs`
и удаляются в teardown). Subject результатов — уникальный на тест (`NatsConnection(result_subject=...)`): общий
`files.tabular_parsed` не трогаем, queue group — боевая `arm-core`. S3 не нужен (`_connect_s3` → None), SMB выключен
(`.env.test`), пул движка БД приложения подменён (lifespan закрывает его последним шагом). Ожидания — с условием и
ограничением по времени, без «голых» пауз.
"""

import asyncio
import importlib
import json
import time
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable
from datetime import UTC, datetime, timedelta
from functools import partial
from typing import Any
from urllib.parse import urlsplit

import httpx
import pytest
import pytest_asyncio
from nats.aio.client import Client as NATS
from nats.aio.msg import Msg
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker

from src.contract import NATS_MSG_ID_HEADER, ack_msg_id, result_msg_id
from src.core.config import get_settings
from src.core.nats import RESULT_QUEUE_GROUP, NatsConnection
from src.models.enums import RecognitionStatus, S3BucketRole
from src.models.files import RecognitionJob, S3Object
from src.services.recognition_exchange import handle_message
from tests.factories import CommittedJobs
from tests.test_recognition_inbox import MESSAGE_VALIDATOR, envelope

pytestmark = pytest.mark.integration

# src/__init__.py реэкспортирует объект app под именем модуля, поэтому модуль берём через importlib
app_module = importlib.import_module("src.app")

READY_TIMEOUT = 5.0
WAIT_TIMEOUT = 10.0
NATS_MONITOR_PORT = 8222  # мониторинг NATS из docs/dev/docker-compose.yml
EXCHANGE_TASKS = 2  # публикатор и вахтёр; S3 не подключено — цикла маппинга нет


# --- инфраструктура ---------------------------------------------------------------------------------------------


async def _wait_until(predicate: Callable[[], Awaitable[bool]], what: str, timeout: float = WAIT_TIMEOUT) -> None:
    deadline = time.monotonic() + timeout
    while not await predicate():
        if time.monotonic() > deadline:
            raise AssertionError(f"за {timeout} с не выполнено: {what}")
        await asyncio.sleep(0.05)


@pytest.fixture
def session_factory(db_engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
    return async_sessionmaker(db_engine, expire_on_commit=False)


@pytest_asyncio.fixture
async def jobs(session_factory: async_sessionmaker[AsyncSession]) -> AsyncIterator[CommittedJobs]:
    committed = CommittedJobs(session_factory)
    try:
        yield committed
    finally:
        await committed.cleanup()


@pytest.fixture
def result_subject() -> str:
    return f"test.result.{uuid.uuid4().hex}"


class _Service:
    """«Сервис распознавания» на голом nats-py: публикует конверты `recognition-result-message` в subject ядра."""

    def __init__(self, result_subject: str) -> None:
        self.result_subject = result_subject
        self.nc = NATS()

    async def start(self) -> None:
        url = get_settings().OPENBAO.NATS.NATS_URL
        try:
            await self.nc.connect(url, connect_timeout=2, allow_reconnect=False)
        except Exception as exc:
            pytest.fail(f"NATS недоступен ({url}, .env.test): {type(exc).__name__}", pytrace=False)

    async def send(self, body: dict[str, Any], msg_id: str) -> None:
        MESSAGE_VALIDATOR.validate(body)
        payload = json.dumps(body, ensure_ascii=False).encode("utf-8")
        await self.nc.publish(self.result_subject, payload, headers={NATS_MSG_ID_HEADER: msg_id})
        await self.nc.flush()

    async def stop(self) -> None:
        await self.nc.close()


@pytest_asyncio.fixture
async def service(result_subject: str) -> AsyncIterator[_Service]:
    svc = _Service(result_subject)
    await svc.start()
    try:
        yield svc
    finally:
        await svc.stop()


async def _row(session_factory: async_sessionmaker[AsyncSession], job_id: uuid.UUID) -> dict[str, Any]:
    async with session_factory() as session:
        row = (await session.execute(select(RecognitionJob.__table__).where(RecognitionJob.id == job_id))).one()
        return dict(row._mapping)


async def _s3_rows(session_factory: async_sessionmaker[AsyncSession], key: str) -> list[dict[str, Any]]:
    async with session_factory() as session:
        rows = await session.execute(select(S3Object.__table__).where(S3Object.key == key))
        return [dict(r._mapping) for r in rows]


async def _status_is(
    session_factory: async_sessionmaker[AsyncSession], job_id: uuid.UUID, status: RecognitionStatus
) -> bool:
    return (await _row(session_factory, job_id))["status"] is status


def _result_key(file_id: uuid.UUID, job_id: uuid.UUID) -> str:
    return f"results/{file_id}/{job_id}/result.json"


async def _queue_subscriptions(subject: str) -> list[dict[str, Any]] | None:
    """Подписки на `subject` по мониторингу NATS (`/connz?subs=detail`); `None` — мониторинг недоступен."""
    host = urlsplit(get_settings().OPENBAO.NATS.NATS_URL).hostname or "localhost"
    url = f"http://{host}:{NATS_MONITOR_PORT}/connz"
    try:
        async with httpx.AsyncClient(timeout=2.0) as client:
            response = await client.get(url, params={"subs": "detail", "limit": 1024})
    except httpx.HTTPError:
        return None
    if response.status_code != httpx.codes.OK:
        return None
    return [
        sub
        for conn in response.json().get("connections", [])
        for sub in conn.get("subscriptions_list_detail") or []
        if sub.get("subject") == subject
    ]


class _EngineStub:
    """Подмена движка БД приложения: lifespan вызывает `dispose()` последним шагом остановки."""

    def __init__(self) -> None:
        self.disposed = 0

    async def dispose(self) -> None:
        self.disposed += 1


# --- (1) реальный lifespan ----------------------------------------------------------------------------------------


async def test_real_lifespan_runs_exchange_over_real_nats(  # noqa: PLR0915 - один сквозной сценарий жизненного цикла
    monkeypatch: pytest.MonkeyPatch,
    session_factory: async_sessionmaker[AsyncSession],
    jobs: CommittedJobs,
    service: _Service,
    result_subject: str,
):
    """NATS_ENABLED=true: `/healthcheck` → `nats=ok`, подписка в queue group `arm-core` активна, публикатор и inbox
    работают с БД; остановка гасит задачи обмена ДО NATS и не оставляет висящих задач."""

    async def no_s3() -> None:
        return None

    engine = _EngineStub()
    monkeypatch.setattr(app_module.settings, "NATS_ENABLED", True)
    monkeypatch.setattr(app_module.settings, "OUTBOX_POLL_SECONDS", 1)  # публикатор не ждёт 5 с между проходами
    monkeypatch.setattr(app_module.settings.OPENBAO.SMB, "SMB_HOST", "")
    monkeypatch.setattr(app_module, "_connect_s3", no_s3)
    monkeypatch.setattr(app_module, "async_session_maker", session_factory)
    monkeypatch.setattr(app_module, "engine", engine)
    # Настоящий `_create_nats` и подписка `_subscribe_results`; отличается только subject результатов
    monkeypatch.setattr(app_module, "NatsConnection", partial(NatsConnection, result_subject=result_subject))

    job_id = await jobs.make_queued_job(reply_deadline=datetime.now(UTC) + timedelta(hours=1))
    file_id = (await jobs.get(job_id)).file_id
    tasks_before = asyncio.all_tasks()
    exchange_done_at_nats_stop: list[bool] = []

    async with app_module.lifespan(app_module.app):
        app = app_module.app
        nats: NatsConnection = app.state.nats
        exchange = app.state.exchange
        assert isinstance(nats, NatsConnection)
        assert nats.result_subject == result_subject
        assert nats.queue == RESULT_QUEUE_GROUP == "arm-core"
        assert exchange is not None
        assert len(exchange.tasks) == EXCHANGE_TASKS

        if not await nats.wait_ready(timeout=READY_TIMEOUT):  # подключение и подписка выполнены
            pytest.fail(f"NATS недоступен ({get_settings().OPENBAO.NATS.NATS_URL}, .env.test)", pytrace=False)
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
            health = (await client.get("/healthcheck")).json()
        assert health["nats"] == "ok"

        subscriptions = await _queue_subscriptions(result_subject)
        if subscriptions is not None:  # мониторинг NATS доступен: подписка ядра видна серверу в queue group
            assert [sub.get("qgroup") for sub in subscriptions] == ["arm-core"]
        assert not any(task.done() for task in exchange.tasks)

        # публикатор lifespan: queued → published
        await _wait_until(
            partial(_status_is, session_factory, job_id, RecognitionStatus.published), "публикация задания"
        )
        # inbox lifespan (подписка на result_subject): ack → processing, результат → done
        await service.send(envelope(job_id, file_id, "processing"), ack_msg_id(job_id, 1))
        await _wait_until(partial(_status_is, session_factory, job_id, RecognitionStatus.processing), "ack")
        await service.send(envelope(job_id, file_id, "done"), result_msg_id(job_id, 1))
        await _wait_until(partial(_status_is, session_factory, job_id, RecognitionStatus.done), "результат")
        assert not any(task.done() for task in exchange.tasks)

        original_stop = nats.stop

        async def recording_stop() -> None:
            exchange_done_at_nats_stop.append(all(task.done() for task in exchange.tasks))
            await original_stop()

        monkeypatch.setattr(nats, "stop", recording_stop)

    assert exchange_done_at_nats_stop == [True]  # задачи обмена остановлены до NATS
    assert all(task.done() for task in exchange.tasks)
    assert not nats.is_connected
    assert engine.disposed == 1

    current = asyncio.current_task()

    async def no_leftover_tasks() -> bool:
        return not {t for t in asyncio.all_tasks() - tasks_before if t is not current and not t.done()}

    await _wait_until(no_leftover_tasks, "все задачи lifespan завершены", timeout=READY_TIMEOUT)
    job = await _row(session_factory, job_id)
    (s3,) = await _s3_rows(session_factory, _result_key(file_id, job_id))
    assert (job["result_s3_object_id"], job["result_msg_id"]) == (s3["id"], f"result:{job_id}:1")
    assert s3["bucket_role"] is S3BucketRole.results


# --- (2) два инстанса ядра в одной queue group ----------------------------------------------------------------------


class _CountingHandler:
    """Обёртка над настоящим `handle_message`: считает вызовы инстанса и незавершённые обработки."""

    def __init__(self, handler: Callable[[Msg], Awaitable[None]]) -> None:
        self._handler = handler
        self.calls = 0
        self.in_flight = 0

    async def __call__(self, msg: Msg) -> None:
        self.calls += 1
        self.in_flight += 1
        try:
            await self._handler(msg)
        finally:
            self.in_flight -= 1


async def test_two_core_instances_handle_result_once(
    session_factory: async_sessionmaker[AsyncSession], jobs: CommittedJobs, service: _Service, result_subject: str
):
    """Два `NatsConnection` в queue group `arm-core` на одном subject: результат обрабатывает ровно один инстанс,
    задание меняется один раз; повторная доставка того же сообщения ничего не меняет."""
    settings = get_settings().OPENBAO.NATS
    handlers = [_CountingHandler(partial(handle_message, session_factory=session_factory)) for _ in range(2)]

    def subscriber(handler: _CountingHandler) -> Callable[[NatsConnection], Awaitable[None]]:
        async def subscribe(conn: NatsConnection) -> None:
            await conn.subscribe_results(handler)

        return subscribe

    cores = [NatsConnection(settings, on_connected=subscriber(h), result_subject=result_subject) for h in handlers]
    job_id = await jobs.make_queued_job(
        status=RecognitionStatus.published,
        published_at=datetime.now(UTC),
        reply_deadline=datetime.now(UTC) + timedelta(hours=1),
    )
    file_id = (await jobs.get(job_id)).file_id
    key = _result_key(file_id, job_id)
    result = envelope(job_id, file_id, "done", metrics={"duration_ms": 7})

    async def settled() -> bool:
        """Сервер доставил всё отправленное обоим инстансам (PING/PONG), и их обработчики завершились."""
        for core in cores:
            await core._client.flush(timeout=2)
        return all(h.in_flight == 0 for h in handlers)

    for core in cores:
        await core.start()
    try:
        for core in cores:
            if not await core.wait_ready(timeout=READY_TIMEOUT):
                pytest.fail(f"NATS недоступен ({settings.NATS_URL}, .env.test)", pytrace=False)
        subscriptions = await _queue_subscriptions(result_subject)
        if subscriptions is not None:
            assert [sub.get("qgroup") for sub in subscriptions] == ["arm-core", "arm-core"]

        await service.send(result, result_msg_id(job_id, 1))
        await _wait_until(partial(_status_is, session_factory, job_id, RecognitionStatus.done), "результат")
        await _wait_until(settled, "обработка доставленных сообщений")

        assert sorted(h.calls for h in handlers) == [0, 1]  # ровно один инстанс получил сообщение
        job = await _row(session_factory, job_id)
        s3_rows = await _s3_rows(session_factory, key)
        assert len(s3_rows) == 1
        assert (job["result_msg_id"], job["result_s3_object_id"]) == (f"result:{job_id}:1", s3_rows[0]["id"])
        assert job["metrics"] == {"duration_ms": 7}

        # повторная доставка того же сообщения (at-least-once): снова один инстанс, без изменений
        await service.send(result, result_msg_id(job_id, 1))

        async def redelivered() -> bool:
            return sum(h.calls for h in handlers) == 2 and await settled()  # noqa: PLR2004 - две доставки

        await _wait_until(redelivered, "повторная доставка обработана")
        assert await _row(session_factory, job_id) == job
        assert await _s3_rows(session_factory, key) == s3_rows
    finally:
        for core in cores:
            await core.stop()