govoplan-core/tests/test_runtime_agents.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace

from govoplan_core.core.runtime_coordination import (
    RuntimeCoordinationError,
    RuntimeIdentity,
    bind_process_runtime_identity,
    process_runtime_identity,
)
from govoplan_core.server.runtime_agent import RuntimeNodeAgent


class _Session:
    def __enter__(self):
        return self

    def __exit__(self, *_args) -> None:
        return None

    def commit(self) -> None:
        return None


class _Database:
    def SessionLocal(self) -> _Session:
        return _Session()


class _Consumer:
    def __init__(self) -> None:
        self.cancelled: list[str] = []
        self.added: list[str] = []

    def cancel_task_queue(self, queue: str) -> None:
        self.cancelled.append(queue)

    def add_task_queue(self, queue: str) -> None:
        self.added.append(queue)


def test_api_runtime_identity_uses_the_installed_distribution_version(
    monkeypatch,
) -> None:
    from govoplan_core.server import default_config

    monkeypatch.setattr(
        default_config,
        "version",
        lambda distribution: "0.1.15"
        if distribution == "govoplan-core"
        else "unexpected",
    )

    assert default_config._core_distribution_version() == "0.1.15"


def test_api_runtime_agent_fails_readiness_on_heartbeat_error() -> None:
    agent = RuntimeNodeAgent(
        settings=SimpleNamespace(
            installation_id="installation-1",
            runtime_role="api",
            runtime_node_id="api-1",
            runtime_heartbeat_seconds=15,
            celery_queues="",
        ),
        software_version="0.1.14",
        module_ids=("core",),
    )
    agent.coordination_healthy = True

    def fail() -> None:
        raise RuntimeCoordinationError("database unavailable")

    agent._heartbeat = fail
    assert asyncio.run(agent._heartbeat_once()) is False
    assert agent.coordination_healthy is False

    agent._heartbeat = lambda: None
    assert asyncio.run(agent._heartbeat_once()) is True
    assert agent.coordination_healthy is True


def test_process_runtime_identity_is_explicit_and_replaceable() -> None:
    from govoplan_core.core import runtime_coordination

    previous = runtime_coordination._process_runtime_identity
    identity = RuntimeIdentity(
        installation_id="installation-1",
        node_id="api-1",
        incarnation="incarnation-1",
        role="api",
        software_version="0.1.14",
        composition_hash="a" * 64,
    )
    try:
        bind_process_runtime_identity(None)
        try:
            process_runtime_identity()
        except RuntimeCoordinationError:
            pass
        else:  # pragma: no cover - assertion branch
            raise AssertionError("An unbound process identity must fail closed")
        bind_process_runtime_identity(identity)
        assert process_runtime_identity() is identity
    finally:
        bind_process_runtime_identity(previous)


def test_worker_disables_consumers_without_reclaiming_stale_identity(
    monkeypatch,
) -> None:
    from govoplan_core import celery_app

    identity = RuntimeIdentity(
        installation_id="installation-1",
        node_id="worker-1",
        incarnation="stale-incarnation",
        role="worker",
        software_version="0.1.14",
        composition_hash="a" * 64,
        queues=("default", "mail"),
    )
    consumer = _Consumer()
    original_state = (
        celery_app._worker_identity,
        celery_app._worker_consumer,
        celery_app._worker_draining,
    )
    celery_app._worker_identity = identity
    celery_app._worker_consumer = consumer
    celery_app._worker_draining = False
    monkeypatch.setattr(celery_app, "get_database", lambda: _Database())

    def reject_heartbeat(*_args, **_kwargs):
        raise RuntimeCoordinationError("stale node incarnation")

    monkeypatch.setattr(
        celery_app,
        "heartbeat_runtime_node",
        reject_heartbeat,
    )
    monkeypatch.setattr(
        celery_app,
        "register_runtime_node",
        lambda *_args, **_kwargs: (_ for _ in ()).throw(
            AssertionError("heartbeat failures must not reclaim an identity")
        ),
    )
    try:
        celery_app._heartbeat_worker_runtime()
        assert consumer.cancelled == ["default", "mail"]
        assert celery_app._worker_draining is True

        monkeypatch.setattr(
            celery_app,
            "heartbeat_runtime_node",
            lambda *_args, **_kwargs: SimpleNamespace(state="active"),
        )
        celery_app._heartbeat_worker_runtime()
        assert consumer.added == ["default", "mail"]
        assert celery_app._worker_draining is False
    finally:
        (
            celery_app._worker_identity,
            celery_app._worker_consumer,
            celery_app._worker_draining,
        ) = original_state


def test_worker_child_replaces_inherited_database_pool(monkeypatch) -> None:
    from govoplan_core import celery_app
    from govoplan_core.core import runtime_coordination

    calls: list[tuple[str, bool]] = []
    previous_process_identity = runtime_coordination._process_runtime_identity
    previous_worker_identity = celery_app._worker_identity
    inherited = RuntimeIdentity(
        installation_id="installation-1",
        node_id="parent-worker",
        incarnation="parent-incarnation",
        role="worker",
        software_version="0.1.14",
        composition_hash="a" * 64,
    )
    monkeypatch.setattr(
        celery_app,
        "configure_database",
        lambda url, *, dispose_previous=False: calls.append(
            (url, dispose_previous)
        ),
    )
    try:
        celery_app._worker_identity = inherited
        bind_process_runtime_identity(inherited)
        celery_app._reset_worker_process_database()

        assert calls == [(celery_app.settings.database_url, True)]
        assert celery_app._worker_identity is None
        try:
            process_runtime_identity()
        except RuntimeCoordinationError:
            pass
        else:  # pragma: no cover - assertion branch
            raise AssertionError("A worker child must discard inherited authority")
    finally:
        celery_app._worker_identity = previous_worker_identity
        bind_process_runtime_identity(previous_process_identity)