126 lines
3.4 KiB
Python
126 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
from govoplan_core.core.runtime_coordination import (
|
|
RuntimeCoordinationError,
|
|
RuntimeIdentity,
|
|
)
|
|
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_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_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
|