"""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, _banned_route_wallet from meshnet_node.server import StubNodeServer from meshnet_contracts import LocalSolanaContracts from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error 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_coverage_endpoint_reports_uncovered_ranges(): """Coverage endpoint returns compressed layer ranges with node counts.""" tracker = TrackerServer(model_presets={ "tiny-model": { "total_layers": 6, "bytes_per_layer": {"bfloat16": 1_000}, }, }) 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", "model": "tiny-model", "shard_start": 0, "shard_end": 2, "hardware_profile": {}, "score": 1.0}, ) coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model") assert coverage_resp["coverage"] == [ {"start_layer": 0, "end_layer": 2, "node_count": 1}, {"start_layer": 3, "end_layer": 5, "node_count": 0}, ] try: _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model") raise AssertionError("Expected 503 for zero-coverage range") except urllib.error.HTTPError as exc: assert exc.code == 503 finally: tracker.stop() def test_tracker_auto_assigns_new_node_to_uncovered_range_first(): """Capability-driven registration fills the first uncovered layer gap.""" tracker = TrackerServer(model_presets={ "tiny-model": { "total_layers": 8, "bytes_per_layer": {"bfloat16": 1_000}, }, }) 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", "model": "tiny-model", "shard_start": 0, "shard_end": 3, "hardware_profile": {}, "score": 1.0}, ) reg = _post_json( f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {"endpoint": "http://127.0.0.1:9002", "model": "tiny-model", "vram_bytes": 10_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) assert reg["directive"]["action"] == "LOAD_SHARD" assert reg["directive"]["start_layer"] == 4 assert reg["directive"]["end_layer"] == 7 finally: tracker.stop() def test_tracker_speed_weighted_vram_assignment_covers_model(): """Three auto-assigned nodes reach 100% coverage with widest range on largest VRAM.""" tracker = TrackerServer(model_presets={ "tiny-model": { "total_layers": 12, "bytes_per_layer": {"bfloat16": 1_000}, }, }) 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", "model": "tiny-model", "vram_bytes": 4_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "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", "model": "tiny-model", "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) _post_json( f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {"endpoint": "http://127.0.0.1:9003", "model": "tiny-model", "vram_bytes": 7_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model") widths = { node["endpoint"]: node["shard_end"] - node["shard_start"] + 1 for node in route_resp["nodes"] } assert widths["http://127.0.0.1:9003"] == 5 assert widths["http://127.0.0.1:9002"] == 4 assert widths["http://127.0.0.1:9001"] == 3 coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model") assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"]) finally: tracker.stop() def test_tracker_speed_is_primary_when_both_nodes_can_cover_gap(): tracker = TrackerServer(model_presets={ "tiny-model": { "total_layers": 12, "bytes_per_layer": {"bfloat16": 1_000}, }, }) tracker_port = tracker.start() try: _post_json( f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {"endpoint": "http://127.0.0.1:9011", "model": "tiny-model", "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) _post_json( f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {"endpoint": "http://127.0.0.1:9012", "model": "tiny-model", "vram_bytes": 15_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0}, ) route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model") widths = { node["endpoint"]: node["shard_end"] - node["shard_start"] + 1 for node in route_resp["nodes"] } assert widths["http://127.0.0.1:9012"] > widths["http://127.0.0.1:9011"] finally: tracker.stop() def test_tracker_registration_directive_is_not_replayed_on_heartbeat(): tracker = TrackerServer(model_presets={ "tiny-model": { "total_layers": 2, "bytes_per_layer": {"bfloat16": 1_000}, }, }) 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:9013", "model": "tiny-model", "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) assert reg["directive"]["action"] == "LOAD_SHARD" hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{reg['node_id']}/heartbeat", {}) assert hb == {} finally: tracker.stop() def test_tracker_reassignment_emits_drop_before_load(): tracker = TrackerServer(model_presets={ "tiny-model": { "total_layers": 4, "bytes_per_layer": {"bfloat16": 1_000}, }, }) tracker_port = tracker.start() try: slow = _post_json( f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {"endpoint": "http://127.0.0.1:9015", "model": "tiny-model", "vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) _post_json( f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {"endpoint": "http://127.0.0.1:9016", "model": "tiny-model", "vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0}, ) hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{slow['node_id']}/heartbeat", {}) assert [directive["action"] for directive in hb["directives"]] == ["DROP_SHARD", "LOAD_SHARD"] finally: tracker.stop() def test_tracker_periodic_rebalance_purges_expired_nodes_without_requests(): tracker = TrackerServer( heartbeat_timeout=0.05, rebalance_interval=0.02, model_presets={"tiny-model": {"total_layers": 1, "bytes_per_layer": {"bfloat16": 1_000}}}, ) 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:9014", "model": "tiny-model", "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) assert reg["node_id"] in tracker._registry time.sleep(0.12) assert reg["node_id"] not in tracker._registry finally: tracker.stop() def test_tracker_faster_node_receives_wider_range_when_capacity_ties(): tracker = TrackerServer(model_presets={ "tiny-model": { "total_layers": 12, "bytes_per_layer": {"bfloat16": 1_000}, }, }) 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", "model": "tiny-model", "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "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", "model": "tiny-model", "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 2.0, "hardware_profile": {}, "score": 1.0}, ) route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model") widths = { node["endpoint"]: node["shard_end"] - node["shard_start"] + 1 for node in route_resp["nodes"] } assert widths["http://127.0.0.1:9002"] > widths["http://127.0.0.1:9001"] finally: tracker.stop() def test_tracker_rebalances_after_middle_range_node_timeout(): """Killing a middle shard queues LOAD_SHARD and restores coverage.""" tracker = TrackerServer( heartbeat_timeout=0.15, model_presets={ "tiny-model": { "total_layers": 12, "bytes_per_layer": {"bfloat16": 1_000}, }, }, ) tracker_port = tracker.start() ids: dict[str, str] = {} try: for port in (9001, 9002, 9003, 9004): reg = _post_json( f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {"endpoint": f"http://127.0.0.1:{port}", "model": "tiny-model", "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, ) ids[str(port)] = reg["node_id"] for node_id in ids.values(): _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {}) time.sleep(0.10) for port in ("9001", "9003", "9004"): _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {}) time.sleep(0.10) coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model") assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"]) heartbeat_responses = [ _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {}) for port in ("9001", "9003", "9004") ] directives = [ directive for response in heartbeat_responses for directive in response.get("directives", []) ] assert any(directive["action"] == "LOAD_SHARD" for directive in directives) 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_tracker_routes_only_nodes_for_requested_model(): """A node registered for one model cannot satisfy another model's route.""" tracker = TrackerServer(model_presets={ "model-a": {"layers_start": 0, "layers_end": 31}, "model-b": {"layers_start": 0, "layers_end": 31}, }) 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, "model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0}, ) try: _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=model-b") raise AssertionError("Expected model-b route to ignore model-a nodes") except urllib.error.HTTPError as exc: assert exc.code == 503 finally: tracker.stop() def test_tracker_assignment_coverage_is_model_scoped(): """Shard assignment gaps are computed only from nodes serving the same model.""" tracker = TrackerServer(model_presets={ "model-a": {"layers_start": 0, "layers_end": 31}, "model-b": {"layers_start": 0, "layers_end": 31}, }) 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, "model": "model-a", "hardware_profile": {}, "wallet_address": "wallet-a", "score": 1.0}, ) assignment = _get_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/assign?model=model-b&device=cpu") assert assignment["shard_start"] == 0 assert assignment["shard_end"] == 15 finally: tracker.stop() def test_tracker_excludes_banned_wallets_from_routes(): """Tracker refuses route candidates whose wallet is banned on-chain.""" contracts = LocalSolanaContracts() contracts.registry.submit_stake("wallet-a", 1_000) contracts.registry.submit_stake("wallet-b", 1_000) contracts.registry.ban_wallet("wallet-b") tracker = TrackerServer(contracts=contracts, minimum_stake=100) 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": {}, "wallet_address": "wallet-a", "score": 1.0}, ) try: _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": {}, "wallet_address": "wallet-b", "score": 1.0}, ) raise AssertionError("Expected 403 for banned wallet registration") except urllib.error.HTTPError as exc: assert exc.code == 403 try: _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=stub-model") raise AssertionError("Expected 503 when only banned wallet covers a layer") except urllib.error.HTTPError as exc: assert exc.code == 503 body = json.loads(exc.read()) assert "layer 16" in body["error"] finally: tracker.stop() def test_tracker_rejects_banned_wallet_registration(): """A banned wallet cannot register with the tracker.""" contracts = LocalSolanaContracts() contracts.registry.submit_stake("wallet-b", 1_000) contracts.registry.ban_wallet("wallet-b") tracker = TrackerServer(contracts=contracts, minimum_stake=100) tracker_port = tracker.start() try: try: _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": 31, "hardware_profile": {}, "wallet_address": "wallet-b", "score": 1.0}, ) raise AssertionError("Expected 403 for banned wallet registration") except urllib.error.HTTPError as exc: assert exc.code == 403 body = json.loads(exc.read()) assert body == {"error": "wallet is banned"} finally: tracker.stop() def test_tracker_ban_registration_guard_reads_contract_state(): """The tracker registration guard reads ban status from the contract facade.""" contracts = LocalSolanaContracts() contracts.registry.ban_wallet("wallet-b") assert _registration_ban_error(contracts, "wallet-b") == "wallet is banned" assert _registration_ban_error(contracts, "wallet-a") is None assert _registration_ban_error(None, "wallet-b") is None def test_gateway_route_guard_rejects_banned_wallet_metadata(): """Gateway rejects tracker routes that include a wallet now banned on-chain.""" contracts = LocalSolanaContracts() contracts.registry.ban_wallet("wallet-b") banned_wallet = _banned_route_wallet( contracts, [{"endpoint": "http://node", "wallet_address": "wallet-b"}], ) assert banned_wallet == "wallet-b" def test_gateway_routes_even_when_compute_node_has_low_stake(): """Routing is not based on a compute node's token balance.""" contracts = LocalSolanaContracts() contracts.registry.submit_stake("wallet-a", 1_000) contracts.registry.submit_stake("wallet-b", 5) 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": {}, "wallet_address": "wallet-a", "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": {}, "wallet_address": "wallet-b", "score": 1.0}, ) gateway = GatewayServer( tracker_url=tracker_url, contracts=contracts, minimum_stake=100, ) 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", "Authorization": "Bearer api-key-a"}, method="POST", ) with urllib.request.urlopen(req) as resp: assert resp.status == 200 assert node_b.received_activations finally: gateway.stop() node_a.stop() node_b.stop() tracker.stop() def test_gateway_records_compute_attribution_after_inference_session(): """Gateway records layer and token attribution for every routed node.""" contracts = LocalSolanaContracts() contracts.registry.submit_stake("wallet-a", 1_000) contracts.registry.submit_stake("wallet-b", 1_000) contracts.payment.fund_api_key("api-key-a", lamports=10_000) tracker = TrackerServer(contracts=contracts, minimum_stake=100) 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": {}, "wallet_address": "wallet-a", "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": {}, "wallet_address": "wallet-b", "score": 1.0}, ) gateway = GatewayServer( tracker_url=tracker_url, contracts=contracts, minimum_stake=100, ) 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", "Authorization": "Bearer api-key-a"}, method="POST", ) with urllib.request.urlopen(req) as resp: assert resp.status == 200 records = contracts.payment.list_attributions() assert [(r.api_key, r.node_wallet, r.layer_start, r.layer_end, r.tokens) for r in records] == [ ("api-key-a", "wallet-a", 0, 15, 1), ("api-key-a", "wallet-b", 16, 31, 1), ] finally: gateway.stop() node_a.stop() node_b.stop() 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()