105 lines
2.9 KiB
Python
105 lines
2.9 KiB
Python
"""Tests for restart recovery of durable Cassandra workers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
from contextlib import nullcontext
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
|
|
MODULE_PATH = (
|
|
Path(__file__).parents[2]
|
|
/ "services"
|
|
/ "hermes"
|
|
/ "scripts"
|
|
/ "recover_cassandra_workers.py"
|
|
)
|
|
|
|
|
|
def _load_module():
|
|
spec = importlib.util.spec_from_file_location("recover_cassandra_workers", MODULE_PATH)
|
|
assert spec and spec.loader
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FakeConnection:
|
|
"""Minimal connection that records closure."""
|
|
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class FakeKanban:
|
|
"""Small fake implementing the Kanban operations used by recovery."""
|
|
|
|
def __init__(self, tasks, *, board_exists: bool = True) -> None:
|
|
self.tasks = tasks
|
|
self.has_board = board_exists
|
|
self.connection = FakeConnection()
|
|
self.reclaimed: list[tuple[str, str]] = []
|
|
self.comments: list[tuple[str, str, str]] = []
|
|
|
|
def board_exists(self, board: str) -> bool:
|
|
assert board == "cassandra"
|
|
return self.has_board
|
|
|
|
def scoped_current_board(self, board: str):
|
|
assert board == "cassandra"
|
|
return nullcontext()
|
|
|
|
def connect(self, *, board: str):
|
|
assert board == "cassandra"
|
|
return self.connection
|
|
|
|
def list_tasks(self, connection):
|
|
assert connection is self.connection
|
|
return self.tasks
|
|
|
|
def reclaim_task(self, connection, task_id: str, *, reason: str) -> bool:
|
|
assert connection is self.connection
|
|
self.reclaimed.append((task_id, reason))
|
|
return task_id != "t_race"
|
|
|
|
def add_comment(self, connection, task_id: str, author: str, body: str) -> None:
|
|
assert connection is self.connection
|
|
self.comments.append((task_id, author, body))
|
|
|
|
|
|
def test_recover_running_tasks_requeues_only_claimed_workers() -> None:
|
|
module = _load_module()
|
|
kanban = FakeKanban(
|
|
[
|
|
SimpleNamespace(id="t_running", status="running"),
|
|
SimpleNamespace(id="t_done", status="done"),
|
|
SimpleNamespace(id="t_race", status="running"),
|
|
]
|
|
)
|
|
|
|
assert module.recover_running_tasks(kanban) == ["t_running"]
|
|
assert kanban.reclaimed == [
|
|
("t_running", module.REASON),
|
|
("t_race", module.REASON),
|
|
]
|
|
assert kanban.comments == [
|
|
(
|
|
"t_running",
|
|
"pod-recovery",
|
|
"Requeued after the Hermes pod restart; the replacement pod will resume this task.",
|
|
)
|
|
]
|
|
assert kanban.connection.closed
|
|
|
|
|
|
def test_recover_running_tasks_is_a_noop_before_board_bootstrap() -> None:
|
|
module = _load_module()
|
|
kanban = FakeKanban([], board_exists=False)
|
|
|
|
assert module.recover_running_tasks(kanban) == []
|
|
assert not kanban.connection.closed
|