Task 11: cluster protocol — subscribe(), ping/pong health, 9 message types

- amqp.py: subscribe() creates exclusive queues bound to routing keys, _rebind_subscriptions() recreates them after reconnect - cluster.py: CLUSTER_NODES, CLUSTER_EVENTS (bounded 1000), CLUSTER_COORDINATOR with auto-promotion, all 6 handlers (register, deregister, pong, event, coordinator_query), ping_node() with 5s timeout + auto-deregister on failure - routers/cluster.py: GET /api/cluster - Wired into app.py lifespan after amqp_connect() - 13 tests covering all handlers, boundaries, and API shape - No passive heartbeats — workers assumed present until ping timeout at work-routing time
2026-07-06 19:22:31 -07:00
parent 454fb3a380
commit 78f7b79494
5 changed files with 575 additions and 1 deletions
@@ -0,0 +1,313 @@
+"""Tests for cluster.py — no live AMQP, all handlers called directly."""
+import asyncio
+from collections import deque
+
+import cluster
+from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM
+
+
+def _reset():
+    cluster.CLUSTER_NODES.clear()
+    cluster.CLUSTER_EVENTS.clear()
+    cluster.CLUSTER_COORDINATOR = None
+    cluster._pending_pings.clear()
+
+
+_published = []
+
+
+async def _fake_publish(exchange, routing_key, payload):
+    _published.append((exchange, routing_key, payload))
+
+
+# ---------- 1. Valid worker registration ----------
+
+
+def test_valid_worker_registration(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
+        {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
+    ))
+
+    assert "jarvis" in cluster.CLUSTER_NODES
+    assert cluster.CLUSTER_NODES["jarvis"]["type"] == "worker"
+    assert cluster.CLUSTER_NODES["jarvis"]["status"] == "active"
+    assert len(cluster.CLUSTER_EVENTS) == 1
+    assert cluster.CLUSTER_EVENTS[0]["category"] == "cluster"
+    assert cluster.CLUSTER_EVENTS[0]["message"] == "Node registered (type=worker)"
+
+    assert len(_published) == 1
+    exchange, rk, payload = _published[0]
+    assert exchange == AMQP_EXCHANGE_ADMIN
+    assert rk == "node.jarvis.admitted"
+    assert payload["type"] == "admitted"
+    assert payload["node_name"] == "jarvis"
+
+
+# ---------- 2. First coordinator auto-promotion ----------
+
+
+def test_first_coordinator_auto_promotion(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.ultron.register",
+        {"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm", "rag"]},
+    ))
+
+    assert cluster.CLUSTER_COORDINATOR == "ultron"
+    assert len(cluster.CLUSTER_EVENTS) == 2
+    assert cluster.CLUSTER_EVENTS[1]["message"] == "Elected as coordinator"
+
+    assert len(_published) == 2
+    # First publish: admitted
+    assert _published[0][1] == "node.ultron.admitted"
+    # Second publish: coord_response
+    exchange, rk, payload = _published[1]
+    assert exchange == AMQP_EXCHANGE_SYSTEM
+    assert rk == "cluster.coordinator.response"
+    assert payload["type"] == "coord_response"
+    assert payload["coordinator"] == "ultron"
+
+
+# ---------- 3. Duplicate node name rejected ----------
+
+
+def test_duplicate_node_rejected(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
+        {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
+    ))
+    _published.clear()
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
+        {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
+    ))
+
+    assert len(cluster.CLUSTER_NODES) == 1
+    assert len(_published) == 1
+    exchange, rk, payload = _published[0]
+    assert exchange == AMQP_EXCHANGE_ADMIN
+    assert rk == "node.jarvis.rejected"
+    assert payload["reason"] == "duplicate_node_name"
+
+
+# ---------- 4. Malformed payload rejected ----------
+
+
+def test_malformed_payload_rejected(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
+        {"node_name": "jarvis"},  # missing node_type and capabilities
+    ))
+
+    assert len(cluster.CLUSTER_NODES) == 0
+    assert len(_published) == 1
+    exchange, rk, payload = _published[0]
+    assert rk == "node.jarvis.rejected"
+    assert payload["reason"] == "malformed_payload"
+
+
+# ---------- 5. Graceful deregistration ----------
+
+
+def test_graceful_deregistration(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
+        {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
+    ))
+    _published.clear()
+
+    asyncio.run(cluster.handle_deregistration(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.deregister",
+        {"node_name": "jarvis"},
+    ))
+
+    assert "jarvis" not in cluster.CLUSTER_NODES
+    assert len(cluster.CLUSTER_EVENTS) == 2
+    assert cluster.CLUSTER_EVENTS[1]["message"] == "Node deregistered"
+
+
+def test_deregister_coordinator_clears(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.ultron.register",
+        {"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm"]},
+    ))
+    assert cluster.CLUSTER_COORDINATOR == "ultron"
+    _published.clear()
+
+    asyncio.run(cluster.handle_deregistration(
+        AMQP_EXCHANGE_ADMIN, "node.ultron.deregister",
+        {"node_name": "ultron"},
+    ))
+
+    assert "ultron" not in cluster.CLUSTER_NODES
+    assert cluster.CLUSTER_COORDINATOR is None
+
+
+# ---------- 6. Pong from known node ----------
+
+
+def test_pong_updates_known_node(monkeypatch):
+    _reset()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.register",
+        {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]},
+    ))
+    original_seen = cluster.CLUSTER_NODES["jarvis"]["last_seen"]
+
+    asyncio.run(cluster.handle_pong(
+        AMQP_EXCHANGE_ADMIN, "node.jarvis.pong",
+        {"node_name": "jarvis", "status": "busy", "load": {"cpu_pct": 80}},
+    ))
+
+    assert cluster.CLUSTER_NODES["jarvis"]["status"] == "busy"
+    assert cluster.CLUSTER_NODES["jarvis"]["load"] == {"cpu_pct": 80}
+    assert cluster.CLUSTER_NODES["jarvis"]["last_seen"] != original_seen
+
+
+# ---------- 7. Pong from unknown node ----------
+
+
+def test_pong_from_unknown_node(caplog, monkeypatch):
+    _reset()
+    caplog.set_level("WARNING")
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_pong(
+        AMQP_EXCHANGE_ADMIN, "node.stranger.pong",
+        {"node_name": "stranger", "status": "active"},
+    ))
+
+    assert "stranger" not in cluster.CLUSTER_NODES
+    assert any("pong from unknown node" in rec.message for rec in caplog.records)
+
+
+# ---------- 8. Event stored in log ----------
+
+
+def test_event_appended_to_log(monkeypatch):
+    _reset()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_event(
+        AMQP_EXCHANGE_SYSTEM, "node.jarvis.event",
+        {"node_name": "jarvis", "severity": "error", "message": "OOM on GPU"},
+    ))
+
+    assert len(cluster.CLUSTER_EVENTS) == 1
+    ev = cluster.CLUSTER_EVENTS[0]
+    assert ev["category"] == "application"
+    assert ev["severity"] == "error"
+    assert ev["node"] == "jarvis"
+
+
+def test_event_log_bounded(monkeypatch):
+    _reset()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    for i in range(1001):
+        cluster._push_event("application", "info", "jarvis", f"event {i}")
+
+    assert len(cluster.CLUSTER_EVENTS) == 1000
+    assert cluster.CLUSTER_EVENTS[0]["message"] == "event 1"
+    assert cluster.CLUSTER_EVENTS[-1]["message"] == "event 1000"
+
+
+# ---------- 9. Coordinator query produces response ----------
+
+
+def test_coordinator_query_response(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_registration(
+        AMQP_EXCHANGE_ADMIN, "node.ultron.register",
+        {"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm"]},
+    ))
+    _published.clear()
+
+    asyncio.run(cluster.handle_coordinator_query(
+        AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.query",
+        {"from": "jarvis", "type": "coord_query"},
+    ))
+
+    assert len(_published) == 1
+    exchange, rk, payload = _published[0]
+    assert exchange == AMQP_EXCHANGE_SYSTEM
+    assert rk == "cluster.coordinator.response"
+    assert payload["coordinator"] == "ultron"
+    assert "nodes" in payload
+
+
+def test_coordinator_query_no_coordinator(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_coordinator_query(
+        AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.query",
+        {"from": "jarvis", "type": "coord_query"},
+    ))
+
+    assert len(_published) == 0
+
+
+# ---------- 10. GET /api/cluster shape ----------
+
+
+def test_cluster_api_shape():
+    _reset()
+    cluster.CLUSTER_NODES["jarvis"] = {
+        "name": "jarvis", "type": "worker", "status": "active",
+        "capabilities": ["llm"], "active_model": None, "load": None,
+        "registered_at": "2026-01-01T00:00:00Z", "last_seen": "2026-01-01T00:00:00Z",
+    }
+    cluster.CLUSTER_COORDINATOR = "ultron"
+    cluster._push_event("cluster", "info", "ultron", "Elected")
+
+    from routers.cluster import cluster_status
+
+    resp = asyncio.run(cluster_status())
+
+    assert "nodes" in resp
+    assert "node_count" in resp
+    assert "coordinator" in resp
+    assert "events" in resp
+    assert resp["node_count"] == 1
+    assert resp["coordinator"] == "ultron"
+    assert "jarvis" in resp["nodes"]
+    assert len(resp["events"]) == 1
+    # No internal keys leaked
+    for node in resp["nodes"].values():
+        for key in node:
+            assert key in {
+                "name", "type", "status", "capabilities", "active_model",
+                "load", "registered_at", "last_seen",
+            }