149 lines
4.5 KiB
Python
149 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
|
|
from rolemesh_gateway.registry import NodeHeartbeat, NodeRegistration, Registry, ServedModel
|
|
|
|
|
|
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"
|
|
|
|
|
|
def test_registry_supports_served_models_and_candidate_selection():
|
|
registry = Registry()
|
|
registry.register(
|
|
NodeRegistration(
|
|
node_id="node-a",
|
|
base_url="http://127.0.0.1:9001",
|
|
served_models=[
|
|
ServedModel(model_id="qwen3-8b", roles=["tutor", "mentor"]),
|
|
ServedModel(model_id="qwen2.5-coder-14b", roles=["code_tutor"]),
|
|
],
|
|
)
|
|
)
|
|
|
|
tutor_candidates = registry.candidates_for_role("tutor")
|
|
assert len(tutor_candidates) == 1
|
|
assert tutor_candidates[0].node_id == "node-a"
|
|
assert tutor_candidates[0].model_id == "qwen3-8b"
|
|
|
|
mentor_candidates = registry.candidates_for_role("mentor")
|
|
assert len(mentor_candidates) == 1
|
|
assert mentor_candidates[0].model_id == "qwen3-8b"
|
|
|
|
code_candidates = registry.candidates_for_role("code_tutor")
|
|
assert len(code_candidates) == 1
|
|
assert code_candidates[0].model_id == "qwen2.5-coder-14b"
|
|
|
|
picked = registry.pick_candidate_for_role("mentor")
|
|
assert picked is not None
|
|
assert picked.model_id == "qwen3-8b"
|
|
|
|
|
|
def test_registry_loads_legacy_roles_as_served_models(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"],
|
|
)
|
|
)
|
|
|
|
reloaded = Registry(persist_path=persist_path)
|
|
candidates = reloaded.candidates_for_role("reviewer")
|
|
assert len(candidates) == 1
|
|
assert candidates[0].model_id == "reviewer"
|