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())