mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor(hub): 用本地 HTTP relay 替代 Worker WebSocket 长连接
Hub 数据面改为 /local/relay/{node_id} HTTP 端点，Worker 通过本机
HTTP 请求转发 tunnel 帧，不再维护 /worker WebSocket 长连接。
Hub 侧:
- 新增 control_plane.rs: Hub 通过 HTTP 回调 Aether app 处理心跳 ACK 和节点状态变更
- 新增 local_relay.rs: 接收本地 HTTP 请求，在 Hub 内部打开 LocalStream 并透传到 proxy
- 移除 worker_conn.rs 及 Worker WebSocket 处理逻辑
- 简化 protocol.rs: 移除 NODE_STATUS 帧类型，抽取通用 encode_frame/decode_payload
Python 侧:
- 删除 tunnel_manager.py 及其 WebSocket 连接管理器 (HubConnectionManager)
- 简化 hub_transport.py 为 HTTP relay 调用
- 新增 src/api/internal/hub.py 接收 Hub 控制面回调 (heartbeat/node-status)
- hub_config.py 移除 WebSocket 相关配置，改为 HTTP relay URL
- service.py 新增 update_tunnel_status 方法
- 删除 src/api/admin/proxy_tunnel.py (旧管理接口)
- proxy_node 缓存 TTL 从 15s 降至 3s 加速状态感知
This commit is contained in:
fawney19
323
tests/e2e_hub_relay.py
Normal file
323
tests/e2e_hub_relay.py
Normal file
@@ -0,0 +1,323 @@
|
||||
"""
|
||||
aether-hub local relay 端到端测试
|
||||
|
||||
测试流程:
|
||||
1. 启动 aether-hub（绑定随机端口）
|
||||
2. 用 websockets 库模拟一个 aether-proxy client 连接到 Hub
|
||||
3. Mock proxy 在收到请求帧后返回固定响应帧
|
||||
4. 通过 Hub 的 /local/relay/{node_id} HTTP API 发送请求
|
||||
5. 验证完整链路: HTTP request -> Hub -> WS frame -> mock proxy -> WS frame -> Hub -> HTTP response
|
||||
|
||||
运行: uv run python tests/e2e_hub_relay.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import gzip
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Protocol constants (mirror aether-hub/src/protocol.rs)
|
||||
# ---------------------------------------------------------------------------
|
||||
HEADER_SIZE = 10
|
||||
|
||||
REQUEST_HEADERS = 0x01
|
||||
REQUEST_BODY = 0x02
|
||||
RESPONSE_HEADERS = 0x03
|
||||
RESPONSE_BODY = 0x04
|
||||
STREAM_END = 0x05
|
||||
STREAM_ERROR = 0x06
|
||||
PING = 0x10
|
||||
PONG = 0x11
|
||||
GOAWAY = 0x12
|
||||
|
||||
FLAG_END_STREAM = 0x01
|
||||
FLAG_GZIP_COMPRESSED = 0x02
|
||||
|
||||
|
||||
def encode_frame(stream_id: int, msg_type: int, flags: int, payload: bytes) -> bytes:
|
||||
header = struct.pack(">I", stream_id) + bytes([msg_type, flags]) + struct.pack(">I", len(payload))
|
||||
return header + payload
|
||||
|
||||
|
||||
def parse_frame(data: bytes) -> tuple[int, int, int, bytes] | None:
|
||||
if len(data) < HEADER_SIZE:
|
||||
return None
|
||||
stream_id = struct.unpack(">I", data[0:4])[0]
|
||||
msg_type = data[4]
|
||||
flags = data[5]
|
||||
payload_len = struct.unpack(">I", data[6:10])[0]
|
||||
if len(data) < HEADER_SIZE + payload_len:
|
||||
return None
|
||||
payload = data[HEADER_SIZE : HEADER_SIZE + payload_len]
|
||||
if flags & FLAG_GZIP_COMPRESSED:
|
||||
payload = gzip.decompress(payload)
|
||||
return stream_id, msg_type, flags, payload
|
||||
|
||||
|
||||
def encode_relay_envelope(meta: dict, body: bytes) -> bytes:
|
||||
meta_json = json.dumps(meta, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
return struct.pack("!I", len(meta_json)) + meta_json + body
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mock aether-proxy: connects to Hub via WebSocket, handles request frames
|
||||
# ---------------------------------------------------------------------------
|
||||
async def mock_proxy(hub_ws_url: str, node_id: str, ready_event: asyncio.Event) -> None:
|
||||
"""Simulate an aether-proxy node that echoes requests as fixed responses."""
|
||||
try:
|
||||
import websockets
|
||||
except ImportError:
|
||||
print("SKIP: websockets package not installed (uv pip install websockets)")
|
||||
sys.exit(1)
|
||||
|
||||
headers = {
|
||||
"X-Node-ID": node_id,
|
||||
"X-Node-Name": f"test-{node_id}",
|
||||
}
|
||||
|
||||
async with websockets.connect(
|
||||
hub_ws_url,
|
||||
additional_headers=headers,
|
||||
max_size=64 * 1024 * 1024,
|
||||
) as ws:
|
||||
ready_event.set()
|
||||
print(f" [mock-proxy] connected to hub as node_id={node_id}")
|
||||
|
||||
request_meta: dict | None = None
|
||||
request_body: bytes = b""
|
||||
|
||||
async for raw_msg in ws:
|
||||
if not isinstance(raw_msg, bytes):
|
||||
continue
|
||||
|
||||
parsed = parse_frame(raw_msg)
|
||||
if parsed is None:
|
||||
continue
|
||||
|
||||
stream_id, msg_type, flags, payload = parsed
|
||||
|
||||
if msg_type == PING:
|
||||
await ws.send(encode_frame(0, PONG, 0, payload))
|
||||
continue
|
||||
|
||||
if msg_type == REQUEST_HEADERS:
|
||||
request_meta = json.loads(payload)
|
||||
print(f" [mock-proxy] stream={stream_id} got REQUEST_HEADERS: {request_meta.get('method')} {request_meta.get('url')}")
|
||||
|
||||
elif msg_type == REQUEST_BODY:
|
||||
request_body = payload
|
||||
is_end = bool(flags & FLAG_END_STREAM)
|
||||
print(f" [mock-proxy] stream={stream_id} got REQUEST_BODY ({len(payload)} bytes, end={is_end})")
|
||||
|
||||
if is_end and request_meta:
|
||||
# Send response: 200 OK with echoed body
|
||||
resp_meta = {
|
||||
"status": 200,
|
||||
"headers": [
|
||||
["content-type", "application/json"],
|
||||
["x-test-echo", "true"],
|
||||
],
|
||||
}
|
||||
resp_meta_json = json.dumps(resp_meta, separators=(",", ":")).encode("utf-8")
|
||||
await ws.send(encode_frame(stream_id, RESPONSE_HEADERS, 0, resp_meta_json))
|
||||
|
||||
echo_body = json.dumps({
|
||||
"echo": True,
|
||||
"received_method": request_meta.get("method"),
|
||||
"received_url": request_meta.get("url"),
|
||||
"received_body_len": len(request_body),
|
||||
}, separators=(",", ":")).encode("utf-8")
|
||||
await ws.send(encode_frame(stream_id, RESPONSE_BODY, 0, echo_body))
|
||||
await ws.send(encode_frame(stream_id, STREAM_END, 0, b""))
|
||||
print(f" [mock-proxy] stream={stream_id} sent response (200, {len(echo_body)} bytes)")
|
||||
|
||||
request_meta = None
|
||||
request_body = b""
|
||||
|
||||
elif msg_type == STREAM_ERROR:
|
||||
error_msg = payload.decode("utf-8", errors="replace")
|
||||
print(f" [mock-proxy] stream={stream_id} got STREAM_ERROR: {error_msg}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test runner
|
||||
# ---------------------------------------------------------------------------
|
||||
async def run_test() -> bool:
|
||||
hub_port = 18085
|
||||
hub_bind = f"127.0.0.1:{hub_port}"
|
||||
hub_binary = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"aether-hub",
|
||||
"target",
|
||||
"release",
|
||||
"aether-hub",
|
||||
)
|
||||
hub_binary = os.path.normpath(hub_binary)
|
||||
|
||||
if not os.path.isfile(hub_binary):
|
||||
print(f"FAIL: hub binary not found at {hub_binary}")
|
||||
print(" run: cd aether-hub && cargo build --release")
|
||||
return False
|
||||
|
||||
# Start aether-hub (with control plane disabled since we don't have the app running)
|
||||
print(f"[1/5] Starting aether-hub on {hub_bind} ...")
|
||||
hub_proc = subprocess.Popen(
|
||||
[
|
||||
hub_binary,
|
||||
"--bind", hub_bind,
|
||||
"--proxy-idle-timeout", "0",
|
||||
"--ping-interval", "30",
|
||||
# Use a non-existent app URL -- control plane callbacks will fail silently
|
||||
"--app-base-url", "http://127.0.0.1:19999",
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
)
|
||||
|
||||
try:
|
||||
# Wait for hub to be ready
|
||||
for _ in range(30):
|
||||
await asyncio.sleep(0.2)
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
resp = await client.get(f"http://{hub_bind}/health", timeout=1.0)
|
||||
if resp.status_code == 200:
|
||||
print(" hub is healthy")
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
else:
|
||||
print("FAIL: hub did not start in time")
|
||||
return False
|
||||
|
||||
# Check initial stats
|
||||
async with httpx.AsyncClient() as client:
|
||||
stats = (await client.get(f"http://{hub_bind}/stats")).json()
|
||||
print(f" initial stats: {stats}")
|
||||
assert stats["proxy_connections"] == 0
|
||||
assert stats["nodes"] == 0
|
||||
|
||||
# Start mock proxy
|
||||
node_id = "test-node-e2e"
|
||||
proxy_ready = asyncio.Event()
|
||||
print(f"\n[2/5] Connecting mock proxy (node_id={node_id}) ...")
|
||||
proxy_task = asyncio.create_task(
|
||||
mock_proxy(f"ws://{hub_bind}/proxy", node_id, proxy_ready)
|
||||
)
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(proxy_ready.wait(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
print("FAIL: mock proxy did not connect in time")
|
||||
return False
|
||||
|
||||
# Give hub a moment to register
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
stats = (await client.get(f"http://{hub_bind}/stats")).json()
|
||||
print(f" stats after connect: {stats}")
|
||||
assert stats["proxy_connections"] == 1, f"expected 1 proxy connection, got {stats['proxy_connections']}"
|
||||
assert stats["nodes"] == 1
|
||||
|
||||
# Send request through local relay
|
||||
print(f"\n[3/5] Sending request via local relay ...")
|
||||
request_body = b'{"model":"test","messages":[]}'
|
||||
envelope = encode_relay_envelope(
|
||||
{
|
||||
"method": "POST",
|
||||
"url": "https://api.example.com/v1/chat/completions",
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
"authorization": "Bearer sk-test-123",
|
||||
},
|
||||
"timeout": 30,
|
||||
},
|
||||
request_body,
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
relay_url = f"http://{hub_bind}/local/relay/{node_id}"
|
||||
resp = await client.post(
|
||||
relay_url,
|
||||
content=envelope,
|
||||
headers={"content-type": "application/vnd.aether.tunnel-envelope"},
|
||||
timeout=10.0,
|
||||
)
|
||||
|
||||
print(f" relay response: status={resp.status_code}")
|
||||
assert resp.status_code == 200, f"expected 200, got {resp.status_code}: {resp.text}"
|
||||
|
||||
echo = resp.json()
|
||||
print(f" echo body: {echo}")
|
||||
assert echo["echo"] is True
|
||||
assert echo["received_method"] == "POST"
|
||||
assert echo["received_url"] == "https://api.example.com/v1/chat/completions"
|
||||
assert echo["received_body_len"] == len(request_body)
|
||||
|
||||
assert resp.headers.get("x-test-echo") == "true"
|
||||
|
||||
# Verify active streams cleaned up
|
||||
print(f"\n[4/5] Verifying stream cleanup ...")
|
||||
await asyncio.sleep(0.2)
|
||||
async with httpx.AsyncClient() as client:
|
||||
stats = (await client.get(f"http://{hub_bind}/stats")).json()
|
||||
print(f" stats after request: {stats}")
|
||||
assert stats["active_streams"] == 0, f"expected 0 active streams, got {stats['active_streams']}"
|
||||
|
||||
# Test error case: request to non-existent node
|
||||
print(f"\n[5/5] Testing error cases ...")
|
||||
async with httpx.AsyncClient() as client:
|
||||
resp = await client.post(
|
||||
f"http://{hub_bind}/local/relay/non-existent-node",
|
||||
content=encode_relay_envelope(
|
||||
{"method": "GET", "url": "https://example.com", "headers": {}, "timeout": 5},
|
||||
b"",
|
||||
),
|
||||
headers={"content-type": "application/vnd.aether.tunnel-envelope"},
|
||||
timeout=5.0,
|
||||
)
|
||||
assert resp.status_code == 503, f"expected 503 for missing node, got {resp.status_code}"
|
||||
assert resp.headers.get("x-aether-tunnel-error") == "connect"
|
||||
print(f" missing node: status={resp.status_code}, error={resp.text}")
|
||||
|
||||
# Test: request from non-loopback should be rejected
|
||||
# (can't easily test from non-loopback, but verify header is present for valid errors)
|
||||
|
||||
# Cleanup: cancel proxy
|
||||
proxy_task.cancel()
|
||||
try:
|
||||
await proxy_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("ALL TESTS PASSED")
|
||||
print("=" * 50)
|
||||
return True
|
||||
|
||||
finally:
|
||||
hub_proc.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
hub_proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
hub_proc.kill()
|
||||
hub_proc.wait()
|
||||
print("\n[cleanup] hub process stopped")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
success = asyncio.run(run_test())
|
||||
sys.exit(0 if success else 1)
|
||||
@@ -1,205 +1,90 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
import struct
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from src.services.proxy_node.hub_config import HubConfig
|
||||
from src.services.proxy_node.hub_transport import HubConnectionManager
|
||||
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
|
||||
from src.services.proxy_node.hub_transport import HubTunnelTransport
|
||||
|
||||
|
||||
class _StubRedis:
|
||||
def __init__(self, *, set_result: bool = True, set_exc: Exception | None = None) -> None:
|
||||
self.set_result = set_result
|
||||
self.set_exc = set_exc
|
||||
self.calls: list[tuple[str, str, int | None, bool | None]] = []
|
||||
class _FakeRelayClient:
|
||||
def __init__(self, response: httpx.Response) -> None:
|
||||
self.response = response
|
||||
self.sent_request: httpx.Request | None = None
|
||||
|
||||
async def set(
|
||||
self,
|
||||
key: str,
|
||||
value: str,
|
||||
ex: int | None = None,
|
||||
nx: bool | None = None,
|
||||
) -> bool:
|
||||
self.calls.append((key, value, ex, nx))
|
||||
if self.set_exc is not None:
|
||||
raise self.set_exc
|
||||
return self.set_result
|
||||
def build_request(self, method: str, url: str, **kwargs: Any) -> httpx.Request:
|
||||
return httpx.Request(method, url, **kwargs)
|
||||
|
||||
async def send(self, request: httpx.Request, *, stream: bool = False) -> httpx.Response:
|
||||
_ = stream
|
||||
self.sent_request = request
|
||||
self.response.request = request
|
||||
return self.response
|
||||
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _build_manager() -> HubConnectionManager:
|
||||
return HubConnectionManager(
|
||||
HubConfig(
|
||||
enabled=True,
|
||||
url="ws://127.0.0.1:8085",
|
||||
connect_timeout_seconds=1.0,
|
||||
ping_interval_seconds=1.0,
|
||||
send_timeout_seconds=1.0,
|
||||
max_streams=16,
|
||||
max_frame_size=1024 * 1024,
|
||||
)
|
||||
def _relay_config() -> HubConfig:
|
||||
return HubConfig(
|
||||
enabled=True,
|
||||
url="http://127.0.0.1:8085",
|
||||
connect_timeout_seconds=1.0,
|
||||
)
|
||||
|
||||
|
||||
def _heartbeat_frame(payload: dict[str, Any]) -> Frame:
|
||||
return Frame(0, MsgType.HEARTBEAT_DATA, 0, json.dumps(payload).encode("utf-8"))
|
||||
@pytest.mark.asyncio
|
||||
async def test_transport_encodes_local_relay_envelope(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
transport = HubTunnelTransport("node-1", timeout=12.0)
|
||||
fake_client = _FakeRelayClient(httpx.Response(200, content=b"ok"))
|
||||
monkeypatch.setattr("src.services.proxy_node.hub_transport.get_hub_config", _relay_config)
|
||||
monkeypatch.setattr(transport, "_relay_client", fake_client)
|
||||
|
||||
request = httpx.Request(
|
||||
"POST",
|
||||
"https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json", "connection": "keep-alive"},
|
||||
content=b'{"hello":"world"}',
|
||||
)
|
||||
|
||||
def _decode_ack(frame: Frame) -> dict[str, Any]:
|
||||
assert frame.msg_type == MsgType.HEARTBEAT_ACK
|
||||
if not frame.payload:
|
||||
return {}
|
||||
return json.loads(frame.payload.decode("utf-8"))
|
||||
response = await transport.handle_async_request(request)
|
||||
assert response.status_code == 200
|
||||
await response.aclose()
|
||||
|
||||
assert fake_client.sent_request is not None
|
||||
payload = fake_client.sent_request.content
|
||||
assert payload is not None
|
||||
meta_len = struct.unpack("!I", payload[:4])[0]
|
||||
meta = json.loads(payload[4 : 4 + meta_len].decode("utf-8"))
|
||||
assert meta == {
|
||||
"method": "POST",
|
||||
"url": "https://example.com/v1/chat/completions",
|
||||
"headers": {"content-type": "application/json"},
|
||||
"timeout": 12,
|
||||
}
|
||||
assert payload[4 + meta_len :] == b'{"hello":"world"}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_heartbeat_normalizes_id_and_updates_db(
|
||||
async def test_transport_maps_relay_timeout_to_read_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = _build_manager()
|
||||
captured_frames: list[Frame] = []
|
||||
heartbeat_calls: list[dict[str, Any]] = []
|
||||
|
||||
async def _fake_send_frame(frame: Frame) -> None:
|
||||
captured_frames.append(frame)
|
||||
|
||||
class _FakeSession:
|
||||
def close(self) -> None:
|
||||
return
|
||||
|
||||
async def _fake_get_redis_client(*, require_redis: bool = False) -> _StubRedis:
|
||||
_ = require_redis
|
||||
return redis
|
||||
|
||||
def _fake_heartbeat(db: Any, **kwargs: Any) -> Any:
|
||||
heartbeat_calls.append(kwargs)
|
||||
return SimpleNamespace(
|
||||
remote_config={"heartbeat_interval": 8, "upgrade_to": "0.2.3"},
|
||||
config_version=5,
|
||||
)
|
||||
|
||||
redis = _StubRedis(set_result=True)
|
||||
monkeypatch.setattr(manager, "_send_frame", _fake_send_frame)
|
||||
monkeypatch.setattr("src.database.create_session", lambda: _FakeSession())
|
||||
monkeypatch.setattr("src.clients.get_redis_client", _fake_get_redis_client)
|
||||
monkeypatch.setattr(
|
||||
"src.services.proxy_node.service.ProxyNodeService.heartbeat", _fake_heartbeat
|
||||
)
|
||||
|
||||
await manager._handle_heartbeat(
|
||||
_heartbeat_frame(
|
||||
{
|
||||
"node_id": "node-1",
|
||||
"heartbeat_session_id": "sess-1",
|
||||
"heartbeat_id": 15.0,
|
||||
"active_connections": 3,
|
||||
"total_requests": 10,
|
||||
"failed_requests": 1,
|
||||
"dns_failures": 2,
|
||||
"stream_errors": 0,
|
||||
"proxy_metadata": {"version": "0.2.1"},
|
||||
}
|
||||
transport = HubTunnelTransport("node-1", timeout=12.0)
|
||||
fake_client = _FakeRelayClient(
|
||||
httpx.Response(
|
||||
504,
|
||||
headers={"x-aether-tunnel-error": "timeout"},
|
||||
content=b"relay timed out",
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr("src.services.proxy_node.hub_transport.get_hub_config", _relay_config)
|
||||
monkeypatch.setattr(transport, "_relay_client", fake_client)
|
||||
|
||||
assert len(redis.calls) == 1
|
||||
assert redis.calls[0][0] == "hub:heartbeat:node-1:sess-1:15"
|
||||
assert heartbeat_calls and heartbeat_calls[0]["node_id"] == "node-1"
|
||||
assert heartbeat_calls[0]["proxy_metadata"] == {"version": "0.2.1"}
|
||||
assert len(captured_frames) == 1
|
||||
request = httpx.Request("GET", "https://example.com")
|
||||
|
||||
ack = _decode_ack(captured_frames[0])
|
||||
assert ack["heartbeat_id"] == 15
|
||||
assert isinstance(ack["heartbeat_id"], int)
|
||||
assert ack["remote_config"] == {"heartbeat_interval": 8, "upgrade_to": "0.2.3"}
|
||||
assert ack["config_version"] == 5
|
||||
assert ack["upgrade_to"] == "0.2.3"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_heartbeat_duplicate_skips_db_update(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
manager = _build_manager()
|
||||
captured_frames: list[Frame] = []
|
||||
heartbeat_called = False
|
||||
|
||||
async def _fake_send_frame(frame: Frame) -> None:
|
||||
captured_frames.append(frame)
|
||||
|
||||
async def _fake_get_redis_client(*, require_redis: bool = False) -> _StubRedis:
|
||||
_ = require_redis
|
||||
return redis
|
||||
|
||||
def _fake_heartbeat(db: Any, **kwargs: Any) -> Any:
|
||||
_ = db, kwargs
|
||||
nonlocal heartbeat_called
|
||||
heartbeat_called = True
|
||||
return SimpleNamespace(remote_config={"heartbeat_interval": 8}, config_version=5)
|
||||
|
||||
redis = _StubRedis(set_result=False)
|
||||
monkeypatch.setattr(manager, "_send_frame", _fake_send_frame)
|
||||
monkeypatch.setattr("src.clients.get_redis_client", _fake_get_redis_client)
|
||||
monkeypatch.setattr(
|
||||
"src.services.proxy_node.service.ProxyNodeService.heartbeat", _fake_heartbeat
|
||||
)
|
||||
|
||||
await manager._handle_heartbeat(
|
||||
_heartbeat_frame({"node_id": "node-1", "heartbeat_id": 77, "total_requests": 20})
|
||||
)
|
||||
|
||||
assert len(redis.calls) == 1
|
||||
assert heartbeat_called is False
|
||||
assert len(captured_frames) == 1
|
||||
ack = _decode_ack(captured_frames[0])
|
||||
assert ack == {"heartbeat_id": 77}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_heartbeat_redis_error_falls_back_to_db_update(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = _build_manager()
|
||||
captured_frames: list[Frame] = []
|
||||
heartbeat_calls: list[dict[str, Any]] = []
|
||||
|
||||
async def _fake_send_frame(frame: Frame) -> None:
|
||||
captured_frames.append(frame)
|
||||
|
||||
class _FakeSession:
|
||||
def close(self) -> None:
|
||||
return
|
||||
|
||||
async def _fake_get_redis_client(*, require_redis: bool = False) -> _StubRedis:
|
||||
_ = require_redis
|
||||
return redis
|
||||
|
||||
def _fake_heartbeat(db: Any, **kwargs: Any) -> Any:
|
||||
heartbeat_calls.append(kwargs)
|
||||
return SimpleNamespace(remote_config=None, config_version=0)
|
||||
|
||||
redis = _StubRedis(set_exc=RuntimeError("redis unavailable"))
|
||||
monkeypatch.setattr(manager, "_send_frame", _fake_send_frame)
|
||||
monkeypatch.setattr("src.database.create_session", lambda: _FakeSession())
|
||||
monkeypatch.setattr("src.clients.get_redis_client", _fake_get_redis_client)
|
||||
monkeypatch.setattr(
|
||||
"src.services.proxy_node.service.ProxyNodeService.heartbeat", _fake_heartbeat
|
||||
)
|
||||
|
||||
await manager._handle_heartbeat(
|
||||
_heartbeat_frame(
|
||||
{
|
||||
"node_id": "node-1",
|
||||
"heartbeat_id": 99,
|
||||
"total_requests": 1,
|
||||
"proxy_metadata": {"version": "0.2.2"},
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert heartbeat_calls and heartbeat_calls[0]["total_requests"] == 1
|
||||
assert heartbeat_calls[0]["proxy_metadata"] == {"version": "0.2.2"}
|
||||
assert len(captured_frames) == 1
|
||||
ack = _decode_ack(captured_frames[0])
|
||||
assert ack == {"heartbeat_id": 99}
|
||||
with pytest.raises(httpx.ReadTimeout, match="relay timed out"):
|
||||
await transport.handle_async_request(request)
|
||||
|
||||
@@ -2,46 +2,40 @@ from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.proxy_node.hub_transport import HubResponseStream
|
||||
from src.services.proxy_node.tunnel_manager import _StreamState
|
||||
from src.services.proxy_node.hub_transport import HubRelayResponseStream
|
||||
|
||||
|
||||
class _Manager:
|
||||
def __init__(self) -> None:
|
||||
self.removed: list[int] = []
|
||||
class _FakeResponse:
|
||||
def __init__(self, chunks: list[bytes]) -> None:
|
||||
self._chunks = chunks
|
||||
self.closed = False
|
||||
|
||||
def remove_stream(self, stream_id: int) -> None:
|
||||
self.removed.append(stream_id)
|
||||
async def aiter_raw(self): # type: ignore[override]
|
||||
for chunk in self._chunks:
|
||||
yield chunk
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hub_response_stream_removes_stream_after_normal_iteration() -> None:
|
||||
manager = _Manager()
|
||||
state = _StreamState(7)
|
||||
state.set_response_headers(200, {})
|
||||
state.push_body_chunk(b"hello")
|
||||
state.set_done()
|
||||
async def test_hub_relay_response_stream_closes_after_iteration() -> None:
|
||||
response = _FakeResponse([b"hello", b"world"])
|
||||
stream = HubRelayResponseStream(response) # type: ignore[arg-type]
|
||||
|
||||
stream = HubResponseStream(manager, state, timeout=0.1)
|
||||
chunks = []
|
||||
async for chunk in stream:
|
||||
chunks.append(chunk)
|
||||
|
||||
assert chunks == [b"hello"]
|
||||
assert manager.removed == [7]
|
||||
assert chunks == [b"hello", b"world"]
|
||||
assert response.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hub_response_stream_removes_stream_after_body_error() -> None:
|
||||
manager = _Manager()
|
||||
state = _StreamState(9)
|
||||
state.set_response_headers(200, {})
|
||||
state.set_error("boom")
|
||||
async def test_hub_relay_response_stream_aclose_closes_response() -> None:
|
||||
response = _FakeResponse([])
|
||||
stream = HubRelayResponseStream(response) # type: ignore[arg-type]
|
||||
|
||||
stream = HubResponseStream(manager, state, timeout=0.1)
|
||||
await stream.aclose()
|
||||
|
||||
with pytest.raises(Exception):
|
||||
async for _chunk in stream:
|
||||
pass
|
||||
|
||||
assert manager.removed == [9]
|
||||
assert response.closed is True
|
||||
|
||||
@@ -1,66 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.proxy_node.hub_config import HubConfig
|
||||
from src.services.proxy_node.hub_transport import HubConnectionManager
|
||||
from src.services.proxy_node.tunnel_manager import TunnelStreamError
|
||||
|
||||
|
||||
def _build_manager() -> HubConnectionManager:
|
||||
return HubConnectionManager(
|
||||
HubConfig(
|
||||
enabled=True,
|
||||
url="ws://127.0.0.1:8085",
|
||||
connect_timeout_seconds=1.0,
|
||||
ping_interval_seconds=1.0,
|
||||
send_timeout_seconds=1.0,
|
||||
max_streams=16,
|
||||
max_frame_size=1024 * 1024,
|
||||
)
|
||||
def test_local_relay_url_uses_http_path() -> None:
|
||||
config = HubConfig(
|
||||
enabled=True,
|
||||
url="http://127.0.0.1:8085",
|
||||
connect_timeout_seconds=1.0,
|
||||
)
|
||||
|
||||
|
||||
def test_record_loop_lag_warning_does_not_degrade() -> None:
|
||||
manager = _build_manager()
|
||||
|
||||
manager._record_loop_lag(1.5)
|
||||
|
||||
assert manager._degraded_until == 0.0
|
||||
|
||||
|
||||
def test_record_loop_lag_degrades_manager(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
manager = _build_manager()
|
||||
now = 1234.0
|
||||
monkeypatch.setattr("src.services.proxy_node.hub_transport._time.monotonic", lambda: now)
|
||||
|
||||
manager._record_loop_lag(4.0)
|
||||
|
||||
assert manager._degraded_until == pytest.approx(now + 12.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_request_rejects_while_manager_degraded(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
manager = _build_manager()
|
||||
manager._degraded_until = time.monotonic() + 5.0
|
||||
|
||||
async def _fake_ensure_connected() -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(manager, "ensure_connected", _fake_ensure_connected)
|
||||
|
||||
with pytest.raises(TunnelStreamError, match="event loop degraded"):
|
||||
await manager.send_request(
|
||||
"node-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=b"{}",
|
||||
timeout=5.0,
|
||||
)
|
||||
|
||||
assert manager._pending_streams == {}
|
||||
assert (
|
||||
config.local_relay_url("node a/1")
|
||||
== "http://127.0.0.1:8085/local/relay/node%20a%2F1"
|
||||
)
|
||||
|
||||
@@ -1,165 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from starlette.websockets import WebSocketState
|
||||
|
||||
from src.services.proxy_node.tunnel_manager import (
|
||||
TunnelConnection,
|
||||
TunnelManager,
|
||||
TunnelStreamError,
|
||||
)
|
||||
from src.services.proxy_node.tunnel_protocol import Frame, MsgType
|
||||
|
||||
|
||||
class _DummyWebSocket:
|
||||
def __init__(self) -> None:
|
||||
self.client_state = WebSocketState.CONNECTED
|
||||
self.sent: list[bytes] = []
|
||||
|
||||
async def send_bytes(self, data: bytes) -> None:
|
||||
self.sent.append(data)
|
||||
|
||||
async def close(self, code: int = 1000, reason: str | None = None) -> None: # noqa: ARG002
|
||||
self.client_state = WebSocketState.DISCONNECTED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_register_and_unregister() -> None:
|
||||
"""register 将连接追加到池中，unregister 按连接实例移除"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
assert manager.connection_count("node-1") == 1
|
||||
assert manager.get_connection("node-1") is conn1
|
||||
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn2)
|
||||
assert manager.connection_count("node-1") == 2
|
||||
|
||||
# unregister conn1 不影响 conn2
|
||||
assert manager.unregister(conn1) is True
|
||||
assert manager.connection_count("node-1") == 1
|
||||
assert manager.get_connection("node-1") is conn2
|
||||
|
||||
# 重复 unregister 返回 False
|
||||
assert manager.unregister(conn1) is False
|
||||
|
||||
# unregister conn2 清空池
|
||||
assert manager.unregister(conn2) is True
|
||||
assert manager.get_connection("node-1") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_least_loaded_selection() -> None:
|
||||
"""get_connection 返回 stream_count 最小的连接"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
manager.register(conn2)
|
||||
|
||||
# 两个都空闲，返回任一（实际返回 min，两者相同时返回第一个）
|
||||
selected = manager.get_connection("node-1")
|
||||
assert selected in (conn1, conn2)
|
||||
|
||||
# 给 conn1 加一个 stream，conn2 应被优先选中
|
||||
conn1.create_stream(2)
|
||||
assert manager.get_connection("node-1") is conn2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dead_connections_cleaned_on_get() -> None:
|
||||
"""get_connection 自动清理 dead 连接"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
manager.register(conn2)
|
||||
|
||||
# 模拟 conn1 断开
|
||||
ws1.client_state = WebSocketState.DISCONNECTED
|
||||
assert manager.get_connection("node-1") is conn2
|
||||
assert manager.connection_count("node-1") == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_removed_connection_frames_ignored() -> None:
|
||||
"""已 unregister 的连接帧不应被处理"""
|
||||
manager = TunnelManager()
|
||||
|
||||
ws1 = _DummyWebSocket()
|
||||
conn1 = TunnelConnection("node-1", "node-1", ws1) # type: ignore[arg-type]
|
||||
ws2 = _DummyWebSocket()
|
||||
conn2 = TunnelConnection("node-1", "node-1", ws2) # type: ignore[arg-type]
|
||||
manager.register(conn1)
|
||||
manager.register(conn2)
|
||||
|
||||
# unregister conn1
|
||||
manager.unregister(conn1)
|
||||
|
||||
ping = Frame(0, MsgType.PING, 0, b"hello")
|
||||
|
||||
# conn1 已不在池中，帧应被忽略
|
||||
await manager.handle_incoming_frame(conn1, ping)
|
||||
# 等待 fire-and-forget task 完成
|
||||
await asyncio.sleep(0.05)
|
||||
assert ws1.sent == []
|
||||
|
||||
# conn2 仍在池中，帧正常处理
|
||||
await manager.handle_incoming_frame(conn2, ping)
|
||||
await asyncio.sleep(0.05)
|
||||
assert len(ws2.sent) == 1
|
||||
pong = Frame.decode(ws2.sent[0])
|
||||
assert pong.msg_type == MsgType.PONG
|
||||
assert pong.payload == b"hello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_streams_from_header() -> None:
|
||||
"""TunnelConnection respects proxy-advertised max_streams (clamped)"""
|
||||
ws = _DummyWebSocket()
|
||||
|
||||
# Explicit value within range
|
||||
conn = TunnelConnection("n", "n", ws, max_streams=256) # type: ignore[arg-type]
|
||||
assert conn.max_streams == 256
|
||||
|
||||
# Clamped to minimum 64
|
||||
conn_low = TunnelConnection("n", "n", ws, max_streams=10) # type: ignore[arg-type]
|
||||
assert conn_low.max_streams == 64
|
||||
|
||||
# Clamped to maximum 2048
|
||||
conn_high = TunnelConnection("n", "n", ws, max_streams=9999) # type: ignore[arg-type]
|
||||
assert conn_high.max_streams == 2048
|
||||
|
||||
# None falls back to TunnelManager.MAX_STREAMS_PER_CONN
|
||||
conn_default = TunnelConnection("n", "n", ws) # type: ignore[arg-type]
|
||||
assert conn_default.max_streams == TunnelManager.MAX_STREAMS_PER_CONN
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_request_respects_per_conn_max_streams() -> None:
|
||||
"""send_request raises TunnelStreamError when per-connection limit is reached"""
|
||||
manager = TunnelManager()
|
||||
ws = _DummyWebSocket()
|
||||
# Set a very low max_streams (clamped to minimum 64)
|
||||
conn = TunnelConnection("node-1", "node-1", ws, max_streams=64) # type: ignore[arg-type]
|
||||
manager.register(conn)
|
||||
|
||||
# Fill up to max_streams
|
||||
for i in range(64):
|
||||
conn.create_stream(i * 2 + 2)
|
||||
|
||||
assert conn.stream_count == 64
|
||||
|
||||
# Next send_request should fail
|
||||
with pytest.raises(TunnelStreamError, match="stream limit reached"):
|
||||
await manager.send_request("node-1", method="GET", url="https://example.com", headers={})
|
||||
Reference in New Issue
Block a user