101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
|
|
from rolemesh_gateway.registry import NodeHeartbeat, NodeRegistration, Registry
|
|
|
|
|
|
def test_registry_persists_round_robin_and_heartbeat_state(tmp_path):
|
|
persist_path = tmp_path / "registry.json"
|
|
registry = Registry(persist_path=persist_path)
|
|
|
|
registry.register(
|
|
NodeRegistration(
|
|
node_id="node-a",
|
|
base_url="http://127.0.0.1:9001",
|
|
roles=["reviewer"],
|
|
)
|
|
)
|
|
registry.register(
|
|
NodeRegistration(
|
|
node_id="node-b",
|
|
base_url="http://127.0.0.1:9002",
|
|
roles=["reviewer"],
|
|
)
|
|
)
|
|
|
|
assert registry.pick_node_for_role("reviewer").node_id == "node-a"
|
|
assert registry.pick_node_for_role("reviewer").node_id == "node-b"
|
|
|
|
node = registry.heartbeat(
|
|
NodeHeartbeat(
|
|
node_id="node-a",
|
|
timestamp=456.0,
|
|
status={"healthy": True},
|
|
metrics=[{"queue_depth": 1}],
|
|
)
|
|
)
|
|
assert node is not None
|
|
assert node.status["metrics"] == [{"queue_depth": 1}]
|
|
|
|
reloaded = Registry(persist_path=persist_path)
|
|
assert reloaded.pick_node_for_role("reviewer").node_id == "node-a"
|
|
assert reloaded.list_nodes()[0].status["timestamp"] == 456.0
|
|
|
|
|
|
def test_registry_heartbeat_returns_none_for_unknown_node():
|
|
registry = Registry()
|
|
assert registry.heartbeat(NodeHeartbeat(node_id="missing")) is None
|
|
|
|
|
|
def test_registry_filters_stale_nodes_from_routing(tmp_path):
|
|
persist_path = tmp_path / "registry.json"
|
|
registry = Registry(persist_path=persist_path, stale_after_s=0.01)
|
|
|
|
fresh = registry.register(
|
|
NodeRegistration(
|
|
node_id="fresh-node",
|
|
base_url="http://127.0.0.1:9001",
|
|
roles=["reviewer"],
|
|
)
|
|
)
|
|
stale = registry.register(
|
|
NodeRegistration(
|
|
node_id="stale-node",
|
|
base_url="http://127.0.0.1:9002",
|
|
roles=["reviewer"],
|
|
)
|
|
)
|
|
|
|
stale.last_seen = stale.last_seen - 60
|
|
registry._save()
|
|
|
|
assert registry.is_stale(stale) is True
|
|
assert registry.is_stale(fresh) is False
|
|
assert [node.node_id for node in registry.nodes_for_role("reviewer", include_stale=False)] == ["fresh-node"]
|
|
assert registry.pick_node_for_role("reviewer").node_id == "fresh-node"
|
|
|
|
|
|
def test_registry_supports_random_selection(monkeypatch):
|
|
registry = Registry()
|
|
registry.register(
|
|
NodeRegistration(
|
|
node_id="node-a",
|
|
base_url="http://127.0.0.1:9001",
|
|
roles=["reviewer"],
|
|
)
|
|
)
|
|
registry.register(
|
|
NodeRegistration(
|
|
node_id="node-b",
|
|
base_url="http://127.0.0.1:9002",
|
|
roles=["reviewer"],
|
|
)
|
|
)
|
|
|
|
monkeypatch.setattr(random, "choice", lambda candidates: candidates[-1])
|
|
|
|
picked = registry.pick_node_for_role("reviewer", strategy="random")
|
|
assert picked is not None
|
|
assert picked.node_id == "node-b"
|