mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-25 21:34:47 -06:00
Patrick Buckley
62ce450b06
Add routing proxy endpoints to the console server:
- POST /v1/api/route/workstreams/new — hash-ring-routed create with
503 retry, target_node pinning, and node_url injection
- POST /v1/api/route/{send,approve,cancel,command,close} — generic
proxy to workstream owner via O(1) bucket lookup
- GET /v1/api/route?ws_id=X — node URL lookup for direct SSE
Wire ConsoleRouter into console lifespan (cache refresh on startup)
and collector discovery loop (version-based cache invalidation).
Add --console-url to channel gateway CLI for multi-node routing.
ChannelRouter routes control-plane through console when set, SSE
connections go direct to server nodes via node_url from create response.
401 lines
14 KiB
Python
401 lines
14 KiB
Python
"""Tests for console routing proxy endpoints (route_create, route_proxy, route_lookup)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
from starlette.testclient import TestClient
|
|
|
|
from turnstone.console.collector import ClusterCollector
|
|
from turnstone.console.router import ConsoleRouter, NodeRef
|
|
from turnstone.core.hash_ring import NoAvailableNodeError
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_mock_collector() -> MagicMock:
|
|
collector = MagicMock(spec=ClusterCollector)
|
|
collector.get_overview.return_value = {
|
|
"nodes": 1,
|
|
"workstreams": 0,
|
|
"states": {"running": 0, "thinking": 0, "attention": 0, "idle": 0, "error": 0},
|
|
"aggregate": {"total_tokens": 0, "total_tool_calls": 0},
|
|
}
|
|
return collector
|
|
|
|
|
|
def _make_mock_router(ready: bool = True) -> MagicMock:
|
|
router = MagicMock(spec=ConsoleRouter)
|
|
router.is_ready.return_value = ready
|
|
router.route.return_value = NodeRef("node-a", "http://a:8080")
|
|
router.generate_ws_id_for_node.return_value = "00ff" + "0" * 28
|
|
return router
|
|
|
|
|
|
def _make_app(
|
|
collector: Any = None,
|
|
router: Any = None,
|
|
) -> Any:
|
|
from turnstone.console.server import _load_static, create_app
|
|
from turnstone.core.auth import AuthConfig
|
|
|
|
_load_static()
|
|
return create_app(
|
|
collector=collector or _make_mock_collector(),
|
|
auth_config=AuthConfig(),
|
|
router=router,
|
|
)
|
|
|
|
|
|
def _make_proxy_post(
|
|
status_code: int = 200,
|
|
json_data: dict[str, Any] | None = None,
|
|
) -> MagicMock:
|
|
"""Create a mock for httpx.AsyncClient.post that returns a fixed response."""
|
|
data = json_data or {"ws_id": "abc123", "name": "test"}
|
|
|
|
async def _mock_post(*args: Any, **kwargs: Any) -> httpx.Response:
|
|
return httpx.Response(
|
|
status_code,
|
|
json=data,
|
|
request=httpx.Request("POST", args[0] if args else "http://test"),
|
|
)
|
|
|
|
mock_post = MagicMock(side_effect=_mock_post)
|
|
return mock_post
|
|
|
|
|
|
def _wire_proxy(app: Any, mock_post: MagicMock | None = None) -> None:
|
|
"""Attach a mock proxy_client to the app (lifespan doesn't run in TestClient)."""
|
|
if mock_post is None:
|
|
mock_post = _make_proxy_post()
|
|
mock_proxy = MagicMock(spec=httpx.AsyncClient)
|
|
mock_proxy.post = mock_post
|
|
app.state.proxy_client = mock_proxy
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests — route_create
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRouteCreate:
|
|
"""POST /v1/api/route/workstreams/new — create via hash-ring routing."""
|
|
|
|
@pytest.fixture()
|
|
def client(self):
|
|
router = _make_mock_router()
|
|
app = _make_app(router=router)
|
|
_wire_proxy(app, _make_proxy_post(json_data={"ws_id": "abc123", "name": "test"}))
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
yield client
|
|
client.close()
|
|
|
|
def test_route_create_proxies_to_node(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"name": "test-ws"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["ws_id"] == "abc123"
|
|
|
|
def test_route_create_injects_node_url(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"name": "test-ws"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["node_url"] == "http://a:8080"
|
|
assert data["node_id"] == "node-a"
|
|
|
|
def test_route_create_resume_ws(self):
|
|
"""resume_ws should route to the node that owns the old workstream."""
|
|
router = _make_mock_router()
|
|
router.route.return_value = NodeRef("node-b", "http://b:8080")
|
|
app = _make_app(router=router)
|
|
_wire_proxy(app, _make_proxy_post(json_data={"ws_id": "old_ws_resumed", "name": "resumed"}))
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
|
|
resp = client.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"resume_ws": "old_ws_id"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["node_url"] == "http://b:8080"
|
|
assert data["node_id"] == "node-b"
|
|
# route() should have been called with the old ws_id
|
|
router.route.assert_called_with("old_ws_id")
|
|
client.close()
|
|
|
|
def test_route_create_target_node(self):
|
|
"""target_node should generate a ws_id that hashes to that node."""
|
|
router = _make_mock_router()
|
|
router.generate_ws_id_for_node.return_value = "00ff" + "0" * 28
|
|
router.route.return_value = NodeRef("node-c", "http://c:8080")
|
|
app = _make_app(router=router)
|
|
_wire_proxy(
|
|
app,
|
|
_make_proxy_post(json_data={"ws_id": "00ff" + "0" * 28, "name": "pinned"}),
|
|
)
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
|
|
resp = client.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"target_node": "node-c"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["node_id"] == "node-c"
|
|
router.generate_ws_id_for_node.assert_called_with("node-c")
|
|
client.close()
|
|
|
|
|
|
class TestRouteCreate503Retry:
|
|
"""503 retry logic in route_create."""
|
|
|
|
def test_route_create_503_retries_on_different_node(self):
|
|
"""If the first node returns 503, retry with a new ws_id targeting a different node."""
|
|
router = _make_mock_router()
|
|
call_count = 0
|
|
|
|
def side_effect_route(ws_id: str) -> NodeRef:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count <= 1:
|
|
# First call returns node-a
|
|
return NodeRef("node-a", "http://a:8080")
|
|
# Subsequent calls return node-b (different node for retry)
|
|
return NodeRef("node-b", "http://b:8080")
|
|
|
|
router.route.side_effect = side_effect_route
|
|
app = _make_app(router=router)
|
|
|
|
post_count = 0
|
|
|
|
async def _mock_post(*args: Any, **kwargs: Any) -> httpx.Response:
|
|
nonlocal post_count
|
|
post_count += 1
|
|
if post_count == 1:
|
|
return httpx.Response(
|
|
503,
|
|
json={"error": "overloaded"},
|
|
request=httpx.Request("POST", args[0] if args else "http://test"),
|
|
)
|
|
return httpx.Response(
|
|
200,
|
|
json={"ws_id": "retry_ws", "name": "retry"},
|
|
request=httpx.Request("POST", args[0] if args else "http://test"),
|
|
)
|
|
|
|
mock_proxy = MagicMock(spec=httpx.AsyncClient)
|
|
mock_proxy.post = MagicMock(side_effect=_mock_post)
|
|
app.state.proxy_client = mock_proxy
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
|
|
resp = client.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"name": "test-ws"},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["ws_id"] == "retry_ws"
|
|
assert data["node_id"] == "node-b"
|
|
assert post_count == 2
|
|
client.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests — route_proxy
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRouteProxy:
|
|
"""POST /v1/api/route/send (and other routed endpoints)."""
|
|
|
|
@pytest.fixture()
|
|
def client(self):
|
|
router = _make_mock_router()
|
|
app = _make_app(router=router)
|
|
_wire_proxy(app, _make_proxy_post(json_data={"status": "ok"}))
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
yield client
|
|
client.close()
|
|
|
|
def test_route_proxy_send(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/send",
|
|
json={"ws_id": "abc123", "message": "hello"},
|
|
)
|
|
assert resp.status_code == 200
|
|
# Verify upstream URL was /v1/api/send (not /v1/api/route/send)
|
|
mock_post = client.app.state.proxy_client.post
|
|
call_args = mock_post.call_args
|
|
assert "/v1/api/send" in call_args[0][0]
|
|
assert "/route/" not in call_args[0][0]
|
|
|
|
def test_route_proxy_approve(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/approve",
|
|
json={"ws_id": "abc123", "approved": True},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_route_proxy_cancel(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/cancel",
|
|
json={"ws_id": "abc123"},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_route_proxy_command(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/command",
|
|
json={"ws_id": "abc123", "command": "status"},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
def test_route_proxy_close(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/workstreams/close",
|
|
json={"ws_id": "abc123"},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests — route_lookup
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRouteLookup:
|
|
"""GET /v1/api/route — look up which node owns a workstream."""
|
|
|
|
@pytest.fixture()
|
|
def client(self):
|
|
router = _make_mock_router()
|
|
app = _make_app(router=router)
|
|
app.state.proxy_client = MagicMock(spec=httpx.AsyncClient)
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
yield client
|
|
client.close()
|
|
|
|
def test_route_lookup(self, client):
|
|
resp = client.get("/v1/api/route?ws_id=abc123")
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["node_url"] == "http://a:8080"
|
|
assert data["node_id"] == "node-a"
|
|
|
|
def test_route_lookup_missing_ws_id(self, client):
|
|
resp = client.get("/v1/api/route")
|
|
assert resp.status_code == 400
|
|
assert "ws_id" in resp.json()["error"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests — not ready / no router -> 503
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRouteNotReady:
|
|
"""When router is None or empty cache, all routing endpoints return 503."""
|
|
|
|
@pytest.fixture()
|
|
def client_no_router(self):
|
|
app = _make_app(router=None)
|
|
app.state.proxy_client = MagicMock(spec=httpx.AsyncClient)
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
yield client
|
|
client.close()
|
|
|
|
@pytest.fixture()
|
|
def client_empty_cache(self):
|
|
router = _make_mock_router(ready=False)
|
|
app = _make_app(router=router)
|
|
app.state.proxy_client = MagicMock(spec=httpx.AsyncClient)
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
yield client
|
|
client.close()
|
|
|
|
def test_route_create_no_router_503(self, client_no_router):
|
|
resp = client_no_router.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"name": "test"},
|
|
)
|
|
assert resp.status_code == 503
|
|
|
|
def test_route_create_empty_cache_503(self, client_empty_cache):
|
|
resp = client_empty_cache.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"name": "test"},
|
|
)
|
|
assert resp.status_code == 503
|
|
|
|
def test_route_proxy_no_router_503(self, client_no_router):
|
|
resp = client_no_router.post(
|
|
"/v1/api/route/send",
|
|
json={"ws_id": "abc", "message": "hello"},
|
|
)
|
|
assert resp.status_code == 503
|
|
|
|
def test_route_lookup_no_router_503(self, client_no_router):
|
|
resp = client_no_router.get("/v1/api/route?ws_id=abc")
|
|
assert resp.status_code == 503
|
|
|
|
def test_route_proxy_empty_cache_503(self, client_empty_cache):
|
|
resp = client_empty_cache.post(
|
|
"/v1/api/route/send",
|
|
json={"ws_id": "abc", "message": "hello"},
|
|
)
|
|
assert resp.status_code == 503
|
|
|
|
def test_route_lookup_empty_cache_503(self, client_empty_cache):
|
|
resp = client_empty_cache.get("/v1/api/route?ws_id=abc")
|
|
assert resp.status_code == 503
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests — NoAvailableNodeError handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRouteNoNode:
|
|
"""When router.route() raises NoAvailableNodeError, endpoints return 503."""
|
|
|
|
@pytest.fixture()
|
|
def client(self):
|
|
router = _make_mock_router()
|
|
router.route.side_effect = NoAvailableNodeError("bucket 0 not assigned")
|
|
app = _make_app(router=router)
|
|
app.state.proxy_client = MagicMock(spec=httpx.AsyncClient)
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
yield client
|
|
client.close()
|
|
|
|
def test_route_create_no_node_503(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/workstreams/new",
|
|
json={"name": "test"},
|
|
)
|
|
assert resp.status_code == 503
|
|
assert "No available node" in resp.json()["error"]
|
|
|
|
def test_route_proxy_no_node_503(self, client):
|
|
resp = client.post(
|
|
"/v1/api/route/send",
|
|
json={"ws_id": "abc", "message": "hello"},
|
|
)
|
|
assert resp.status_code == 503
|
|
|
|
def test_route_lookup_no_node_503(self, client):
|
|
resp = client.get("/v1/api/route?ws_id=abc")
|
|
assert resp.status_code == 503
|