RoleMesh-Gateway/tests/test_gateway.py

from __future__ import annotations

import asyncio
from pathlib import Path

import httpx
import yaml


def _write_gateway_config(path: Path) -> Path:
    path.write_text(
        yaml.safe_dump(
            {
                "version": 1,
                "default_model": "writer",
                "auth": {
                    "client_api_keys": ["client-secret"],
                    "node_api_keys": ["node-secret"],
                },
                "models": {
                    "writer": {
                        "type": "proxy",
                        "openai_model_name": "writer",
                        "proxy_url": "http://127.0.0.1:8012",
                        "defaults": {"temperature": 0.6},
                    },
                    "reviewer": {
                        "type": "discovered",
                        "openai_model_name": "reviewer",
                        "role": "reviewer",
                        "strategy": "round_robin",
                    },
                },
            }
        )
    )
    return path


def _create_gateway_app(tmp_path: Path):
    from rolemesh_gateway.main import create_app

    return create_app(
        _write_gateway_config(tmp_path / "models.yaml"),
        registry_path=tmp_path / "registry.json",
    )


async def _request(app, method: str, path: str, **kwargs) -> httpx.Response:
    transport = httpx.ASGITransport(app=app)
    async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
        return await client.request(method, path, **kwargs)


def test_models_requires_client_auth(tmp_path):
    app = _create_gateway_app(tmp_path)
    unauthorized = asyncio.run(_request(app, "GET", "/v1/models"))
    assert unauthorized.status_code == 401

    authorized = asyncio.run(
        _request(app, "GET", "/v1/models", headers={"x-api-key": "client-secret"})
    )
    assert authorized.status_code == 200
    body = authorized.json()
    assert {item["id"] for item in body["data"]} == {"writer", "reviewer"}
    asyncio.run(app.state.upstream.close())


def test_chat_completions_applies_defaults_for_proxy_model(tmp_path):
    app = _create_gateway_app(tmp_path)
    calls = {}

    async def fake_chat(base_url, payload):
        calls["base_url"] = base_url
        calls["payload"] = payload
        return {"id": "cmpl-1", "choices": [{"message": {"role": "assistant", "content": "ok"}}]}

    app.state.upstream.chat_completions = fake_chat

    response = asyncio.run(
        _request(
            app,
            "POST",
            "/v1/chat/completions",
            headers={"x-api-key": "client-secret"},
            json={
                "model": "writer",
                "messages": [{"role": "user", "content": "hello"}],
            },
        )
    )

    assert response.status_code == 200
    assert calls["base_url"] == "http://127.0.0.1:8012"
    assert calls["payload"]["temperature"] == 0.6
    assert response.json()["choices"][0]["message"]["content"] == "ok"
    asyncio.run(app.state.upstream.close())


def test_discovered_model_without_registered_node_returns_503(tmp_path):
    app = _create_gateway_app(tmp_path)
    response = asyncio.run(
        _request(
            app,
            "POST",
            "/v1/chat/completions",
            headers={"x-api-key": "client-secret"},
            json={
                "model": "reviewer",
                "messages": [{"role": "user", "content": "hello"}],
            },
        )
    )

    assert response.status_code == 503
    assert response.json()["error"]["code"] == "no_upstream"
    asyncio.run(app.state.upstream.close())


def test_node_registration_and_heartbeat_update_registry_state(tmp_path):
    app = _create_gateway_app(tmp_path)
    register = asyncio.run(
        _request(
            app,
            "POST",
            "/v1/nodes/register",
            headers={"x-rolemesh-node-key": "node-secret"},
            json={
                "node_id": "node-a",
                "base_url": "http://127.0.0.1:9001",
                "roles": ["reviewer"],
                "meta": {"gpu": "test"},
            },
        )
    )
    assert register.status_code == 200

    heartbeat = asyncio.run(
        _request(
            app,
            "POST",
            "/v1/nodes/heartbeat",
            headers={"x-rolemesh-node-key": "node-secret"},
            json={
                "node_id": "node-a",
                "timestamp": 123.0,
                "status": {"healthy": True},
                "metrics": [{"device": {"id": "gpu:0"}}],
            },
        )
    )

    assert heartbeat.status_code == 200
    node = heartbeat.json()["node"]
    assert node["status"]["healthy"] is True
    assert node["status"]["timestamp"] == 123.0
    assert node["status"]["metrics"] == [{"device": {"id": "gpu:0"}}]
    asyncio.run(app.state.upstream.close())