feat: add tracker registration routing
This commit is contained in:
Dobromir Popov
365
tests/test_tracker_routing.py
Normal file
365
tests/test_tracker_routing.py
Normal file
@@ -0,0 +1,365 @@
|
||||
"""US-003 tests: tracker node registration and route selection."""
|
||||
|
||||
import http.server
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
from meshnet_gateway.server import GatewayServer
|
||||
from meshnet_node.server import StubNodeServer
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
|
||||
|
||||
def _post_json(url: str, payload: dict) -> dict:
|
||||
data = json.dumps(payload).encode()
|
||||
req = urllib.request.Request(
|
||||
url, data=data, headers={"Content-Type": "application/json"}, method="POST"
|
||||
)
|
||||
with urllib.request.urlopen(req) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _get_json(url: str) -> dict:
|
||||
with urllib.request.urlopen(url) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def test_tracker_node_registration():
|
||||
"""A node can register with the tracker and receives a node_id."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
resp = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{
|
||||
"endpoint": "http://127.0.0.1:9000",
|
||||
"shard_start": 0,
|
||||
"shard_end": 15,
|
||||
"hardware_profile": {"gpu": "rtx3090"},
|
||||
"score": 1.0,
|
||||
},
|
||||
)
|
||||
assert "node_id" in resp
|
||||
assert isinstance(resp["node_id"], str)
|
||||
assert len(resp["node_id"]) > 0
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_route_selection():
|
||||
"""Tracker returns ordered route when nodes collectively cover all layers."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9002", "shard_start": 16, "shard_end": 31,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
route_resp = _get_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
|
||||
)
|
||||
assert "route" in route_resp
|
||||
assert route_resp["route"] == [
|
||||
"http://127.0.0.1:9001",
|
||||
"http://127.0.0.1:9002",
|
||||
]
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_route_error_no_coverage():
|
||||
"""Tracker returns 503 when registered nodes do not cover all required layers."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
# Only half the layers covered — route must fail.
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 15,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
try:
|
||||
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
||||
raise AssertionError("Expected 503 but request succeeded")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 503
|
||||
body = json.loads(exc.read())
|
||||
assert "error" in body
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_route_error_no_nodes():
|
||||
"""Tracker returns 503 with clear error when the registry is empty."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
try:
|
||||
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
||||
raise AssertionError("Expected 503 but request succeeded")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 503
|
||||
body = json.loads(exc.read())
|
||||
assert "error" in body
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_heartbeat_updates_node():
|
||||
"""Sending a heartbeat for a registered node succeeds."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
reg = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
node_id = reg["node_id"]
|
||||
|
||||
hb_resp = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {}
|
||||
)
|
||||
assert hb_resp == {}
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_heartbeat_expiry():
|
||||
"""Nodes that miss their heartbeat window are excluded from routes."""
|
||||
tracker = TrackerServer(heartbeat_timeout=0.05) # 50 ms
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
# Immediately after registration the node is alive.
|
||||
route_resp = _get_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model"
|
||||
)
|
||||
assert route_resp["route"] == ["http://127.0.0.1:9001"]
|
||||
|
||||
# Let the heartbeat window expire.
|
||||
time.sleep(0.15)
|
||||
|
||||
try:
|
||||
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
||||
raise AssertionError("Expected 503 after heartbeat expiry")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 503
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_heartbeat_expiry_removes_node_from_registry():
|
||||
"""Expired nodes are removed, not merely hidden from route responses."""
|
||||
tracker = TrackerServer(heartbeat_timeout=0.05)
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
reg = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
node_id = reg["node_id"]
|
||||
time.sleep(0.15)
|
||||
|
||||
try:
|
||||
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
||||
raise AssertionError("Expected 503 after heartbeat expiry")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 503
|
||||
|
||||
try:
|
||||
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {})
|
||||
raise AssertionError("Expected 404 for removed node")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 404
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_route_rejects_non_extending_overlap():
|
||||
"""Overlapping shards that do not extend coverage cannot form a complete route."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 10,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9002", "shard_start": 0, "shard_end": 5,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
try:
|
||||
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model")
|
||||
raise AssertionError("Expected 503 for incomplete overlapping route")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 503
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_registration_rejects_invalid_payload():
|
||||
"""Registration errors return a defined JSON 400 response."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "not-a-url", "shard_start": 10, "shard_end": 0,
|
||||
"hardware_profile": []},
|
||||
)
|
||||
raise AssertionError("Expected 400 for invalid registration")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 400
|
||||
body = json.loads(exc.read())
|
||||
assert "error" in body
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_gateway_returns_json_503_when_tracker_unavailable():
|
||||
"""Tracker-backed gateway failures are reported as JSON HTTP errors."""
|
||||
gateway = GatewayServer(tracker_url="http://127.0.0.1:9")
|
||||
gateway_port = gateway.start()
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "stub-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
urllib.request.urlopen(req, timeout=5)
|
||||
raise AssertionError("Expected 503 for unavailable tracker")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 503
|
||||
assert exc.headers["Content-Type"] == "application/json"
|
||||
body = json.loads(exc.read())
|
||||
assert body == {"error": "tracker unavailable"}
|
||||
finally:
|
||||
gateway.stop()
|
||||
|
||||
|
||||
def test_gateway_returns_json_503_for_malformed_tracker_route():
|
||||
"""Malformed tracker route payloads do not drop the gateway connection."""
|
||||
|
||||
class MalformedRouteHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args):
|
||||
pass
|
||||
|
||||
def do_GET(self):
|
||||
body = json.dumps({"route": []}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
fake_tracker = http.server.HTTPServer(("127.0.0.1", 0), MalformedRouteHandler)
|
||||
fake_thread = threading.Thread(target=fake_tracker.serve_forever, daemon=True)
|
||||
fake_thread.start()
|
||||
gateway = GatewayServer(tracker_url=f"http://127.0.0.1:{fake_tracker.server_address[1]}")
|
||||
gateway_port = gateway.start()
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "stub-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
urllib.request.urlopen(req, timeout=5)
|
||||
raise AssertionError("Expected 503 for malformed tracker route")
|
||||
except urllib.error.HTTPError as exc:
|
||||
assert exc.code == 503
|
||||
assert exc.headers["Content-Type"] == "application/json"
|
||||
body = json.loads(exc.read())
|
||||
assert body == {"error": "tracker unavailable"}
|
||||
finally:
|
||||
gateway.stop()
|
||||
fake_tracker.shutdown()
|
||||
fake_tracker.server_close()
|
||||
fake_thread.join(timeout=1)
|
||||
|
||||
|
||||
def test_two_node_pipeline_via_tracker():
|
||||
"""End-to-end: nodes register with tracker; gateway discovers route dynamically."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
|
||||
node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
|
||||
port_a = node_a.start()
|
||||
|
||||
node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
|
||||
port_b = node_b.start()
|
||||
|
||||
_post_json(
|
||||
f"{tracker_url}/v1/nodes/register",
|
||||
{"endpoint": f"http://127.0.0.1:{port_a}", "shard_start": 0, "shard_end": 15,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
_post_json(
|
||||
f"{tracker_url}/v1/nodes/register",
|
||||
{"endpoint": f"http://127.0.0.1:{port_b}", "shard_start": 16, "shard_end": 31,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
gateway = GatewayServer(tracker_url=tracker_url)
|
||||
gateway_port = gateway.start()
|
||||
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "stub-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}).encode()
|
||||
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{gateway_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
|
||||
with urllib.request.urlopen(req) as resp:
|
||||
assert resp.status == 200
|
||||
body = json.loads(resp.read())
|
||||
|
||||
assert body["id"] == "chatcmpl-stub"
|
||||
assert body["object"] == "chat.completion"
|
||||
assert body["choices"][0]["message"]["content"] == "stub response to: hello"
|
||||
assert node_b.received_activations, "node B did not receive activations from node A"
|
||||
assert not node_a.received_activations, "node A should not receive activations"
|
||||
finally:
|
||||
gateway.stop()
|
||||
node_a.stop()
|
||||
node_b.stop()
|
||||
tracker.stop()
|
||||
Reference in New Issue
Block a user