cAIc/tests/test_cluster.py

"""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",
            }