#!/usr/bin/env python3 """ scripts/test_ws_state.py — the push cache, against a real WebSocket server. venv/bin/python scripts/test_ws_state.py Stands up a local mock Solana RPC WebSocket that speaks slotSubscribe, accountSubscribe and accountNotification, then drives modules/ws_state.py against it end to end. No network, no key, nothing signed. The two things worth proving, because both are silent when wrong: 1. A subscription id is NOT a request id. accountNotification carries the SERVER's `params.subscription`, and binding notifications to our own outgoing request id would route every one of them to the wrong cache entry — or to none — while looking perfectly healthy from outside. 2. A dead socket must cost nothing. If the WS never connects, drops mid-run, or is switched off, every consumer has to behave exactly as it did before this module existed. This is an optimisation, and an optimisation that can break a send is not one. """ from __future__ import annotations import asyncio import json import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) FAILURES: list[str] = [] CHECKS = 0 def check(label: str, condition: bool, detail: str = "") -> None: global CHECKS CHECKS += 1 if condition: print(f" \033[32m✓\033[0m {label}") else: FAILURES.append(f"{label}{f' — {detail}' if detail else ''}") print(f" \033[31m✗\033[0m {label}" + (f" — {detail}" if detail else "")) ALT_ADDRESS = "9WzDXwBbmkg8ZTbNMqUxvQRAyrZzDsGYdLVL9zYtAWWM" # Deliberately far from any request id we will issue, so a test that passes # cannot be passing by coincidence of the two number spaces overlapping. SERVER_SUB_ID = 77_001 class MockServer: """Minimal Solana-shaped WS: confirms subscriptions, emits notifications.""" def __init__(self) -> None: self.subscribed: list[str] = [] self.slot_subs = 0 self._ws = None async def handler(self, ws) -> None: self._ws = ws async for raw in ws: msg = json.loads(raw) method, req_id = msg.get("method"), msg.get("id") if method == "slotSubscribe": self.slot_subs += 1 await ws.send(json.dumps( {"jsonrpc": "2.0", "result": 1, "id": req_id})) elif method == "accountSubscribe": self.subscribed.append(msg["params"][0]) await ws.send(json.dumps( {"jsonrpc": "2.0", "result": SERVER_SUB_ID, "id": req_id})) async def emit_slot(self, slot: int = 1) -> None: await self._ws.send(json.dumps({ "jsonrpc": "2.0", "method": "slotNotification", "params": {"result": {"slot": slot}, "subscription": 1}, })) async def emit_account_change(self) -> None: await self._ws.send(json.dumps({ "jsonrpc": "2.0", "method": "accountNotification", "params": {"result": {"value": {}}, "subscription": SERVER_SUB_ID}, })) async def run() -> None: import websockets from modules.hot_state import get_hot_state from modules.ws_state import WsState hot = get_hot_state() server = MockServer() async with websockets.serve(server.handler, "127.0.0.1", 0) as srv: port = srv.sockets[0].getsockname()[1] url = f"ws://127.0.0.1:{port}" ws_state = WsState(ws_url=url) ws_state.note_alt_addresses([ALT_ADDRESS]) ws_state.start() # Let it connect and subscribe. for _ in range(50): await asyncio.sleep(0.05) if ws_state.status()["connected"] and server.subscribed: break print("\n\033[1m1. it connects and subscribes\033[0m") check("connected to the mock endpoint", ws_state.status()["connected"]) check("subscribed to slots", server.slot_subs == 1) check("subscribed to the ALT the send path used", server.subscribed == [ALT_ADDRESS], f"got {server.subscribed}") print("\n\033[1m2. a subscription id is not a request id\033[0m") subs = ws_state._subs check("the SERVER's subscription id is what got bound", SERVER_SUB_ID in subs, f"bound ids are {sorted(subs)} — notifications would not route") check("it maps back to the right cache entry", subs.get(SERVER_SUB_ID) == ("alt", ALT_ADDRESS)) check("the request id was not left bound", not ws_state._pending_subs) print("\n\033[1m3. an on-chain change drops the stale cache entry\033[0m") hot.put("alt", ALT_ADDRESS, ["stale-value"]) value, age = hot.peek("alt", ALT_ADDRESS) check("a value can be pushed in", value == ["stale-value"] and age < 5.0) await server.emit_account_change() for _ in range(40): await asyncio.sleep(0.05) if hot.peek("alt", ALT_ADDRESS)[0] is None: break value, _ = hot.peek("alt", ALT_ADDRESS) check("the changed account was invalidated", value is None, "a cached ALT the chain has changed is WRONG, not merely old") print("\n\033[1m4. slot notifications drive blockhash refresh\033[0m") before = ws_state.status()["slots_seen"] await server.emit_slot(2) for _ in range(40): await asyncio.sleep(0.05) if ws_state.status()["slots_seen"] > before: break check("slot notifications are observed", ws_state.status()["slots_seen"] > before) # No rpc_url was configured, so the refresh must decline quietly # rather than raising into the socket loop and killing it. check("still connected after a refresh it could not perform", ws_state.status()["connected"], "a failed push must never take down the subscription") await ws_state.stop() print("\n\033[1m5. it is never required\033[0m") from modules.ws_state import WsState as W dead = W(ws_url="ws://127.0.0.1:1") # nothing listening dead.start() await asyncio.sleep(0.4) check("an unreachable endpoint does not raise", True) check("it reports itself as not connected", not dead.status()["connected"]) await dead.stop() no_url = W(ws_url="") no_url.start() check("no URL configured is a no-op, not an error", not no_url.status()["connected"]) await no_url.stop() print("\n\033[1m6. hot_state keeps its meaning\033[0m") from modules.hot_state import HotState h = HotState() h.put("blockhash", "u", "abc") st = h.status() entry = st["entries"]["blockhash:u"] check("a push counts as a refresh, not a hit", entry["refreshes"] == 1 and entry["hits"] == 0, "counting pushes as hits would inflate hit_rate with values " "no trade ever asked for") check("hit_rate stays honest with only pushes", st["hit_rate"] is None) check("peek on an unknown key returns (None, inf)", h.peek("alt", "nope") == (None, float("inf"))) def test_ws_url_override() -> None: """SOLANA_WS_STATE_URL was documented from day one and never read.""" import os from modules.ws_state import WsState print("\n\033[1m7. the documented endpoint override actually works\033[0m") os.environ["SOLANA_WS_STATE_URL"] = "ws://override.example:9999" try: w = WsState() # start() resolves the URL; it will fail to connect, which is fine — # the claim under test is which URL it CHOSE, not whether it dialled. w.start(ws_url="ws://from-rpc-primary.example:1111") check( "the override beats SOLANA_RPC_WS_PRIMARY", w._ws_url == "ws://override.example:9999", f"chose {w._ws_url!r} — a documented knob that does nothing is " f"worse than an undocumented one, because it is trusted", ) finally: os.environ.pop("SOLANA_WS_STATE_URL", None) w2 = WsState() w2.start(ws_url="ws://from-rpc-primary.example:1111") check("without it, the RPC primary is still used", w2._ws_url == "ws://from-rpc-primary.example:1111") def test_grpc_preflight() -> None: import os from modules import grpc_preflight as g print("\n\033[1m8. gRPC preflight names the missing piece\033[0m") for key in ("SOLANA_GRPC_ENDPOINT", "SOLANA_GRPC_X_TOKEN"): os.environ.pop(key, None) st = g.status() check("unconfigured is not 'ready'", not st["ready"]) # Deliberately NOT asserting the exact wording. render() says "nothing # installed for it" only when grpcio is also absent, so pinning that # string makes the test pass or fail on whether the machine running it # happens to have grpcio — which is not a property of this code. The # invariant that matters is that an unconfigured box is never told it is # ready. check("and never claims to be ready", "ready —" not in g.render()) check("configured is False with no endpoint", not st["configured"]) # The scheme traps. Each must be caught with its OWN reason — "invalid" # would send someone to check the network for a typo. for bad, want in (("wss://n.rpcpool.com", "WebSocket"), ("https://n.rpcpool.com", "drop the scheme")): os.environ["SOLANA_GRPC_ENDPOINT"] = bad c = [c for c in g.status()["checks"] if "ENDPOINT" in c["name"]][0] check(f"{bad} is rejected with the real reason", not c["ok"] and want in c["fix"]) os.environ["SOLANA_GRPC_ENDPOINT"] = "n.rpcpool.com:443" c = [c for c in g.status()["checks"] if "ENDPOINT" in c["name"]][0] check("host:port is accepted", c["ok"]) os.environ["SOLANA_GRPC_X_TOKEN"] = "t" st = g.status() check("a missing token is still reported when the endpoint is fine", all(c["ok"] for c in st["checks"] if "TOKEN" in c["name"])) check("no token in the reported host", "t" not in st["endpoint_host"] or "rpcpool" in st["endpoint_host"]) for key in ("SOLANA_GRPC_ENDPOINT", "SOLANA_GRPC_X_TOKEN"): os.environ.pop(key, None) def test_grpc_blockhash_shape() -> None: """The push must be indistinguishable from the fetch it replaces.""" print("\n\033[1m9. gRPC pushes the same type the send path fetches\033[0m") from solders.hash import Hash from modules.grpc_client import GrpcState BH = "EETubP5AKHgjPAhzPAFcb8BAY1hMH639CWCFTqi3hq1k" class FakeResp: slot = 1 blockhash = BH last_valid_block_height = 9 pushed = GrpcState._as_rpc_shape(FakeResp()) # solana_executor._get_latest_blockhash returns Hash.from_string(...). # THE bug this guards: an earlier version built a GetLatestBlockhashResp # because that is what the JSON-RPC method is called, which would have # put the wrong type into the cache the transaction compiler reads — # every trade failing at signing, on a box just "upgraded" to a faster # feed. fetched = Hash.from_string(BH) check("the pushed value is a solders Hash", isinstance(pushed, Hash), f"got {type(pushed).__name__} — the send path compiles with this") check("it equals what the RPC fetch would return", pushed == fetched) check("it is NOT a response wrapper", type(pushed).__name__ == "Hash", "GetLatestBlockhashResp would poison the cache") def test_grpc_never_required() -> None: print("\n\033[1m10. gRPC is never required\033[0m") import os from modules.grpc_client import GrpcState for key in ("SOLANA_GRPC_ENDPOINT", "SOLANA_GRPC_X_TOKEN"): os.environ.pop(key, None) g = GrpcState() g.start() # no endpoint at all check("no endpoint is a silent no-op", not g.status()["connected"]) check("and it reports itself honestly", g.status()["endpoint_host"] == "") # Configured but with the stubs/deps question unresolved: start() must # still return quietly rather than raising or spinning on reconnects. os.environ["SOLANA_GRPC_ENDPOINT"] = "unreachable.invalid:443" g2 = GrpcState() g2.start() check("a configured-but-unusable endpoint does not raise", True) os.environ.pop("SOLANA_GRPC_ENDPOINT", None) def test_liveness_not_just_connected() -> None: """CONNECTED is not LIVE, and the difference costs a cold fetch.""" print("\n\033[1mliveness: a socket being open proves nothing\033[0m") import time as _t from modules.ws_state import _STALE_AFTER_SECS, WsState w = WsState(ws_url="wss://x/y", rpc_url="https://x/y") check("a fresh object is not live", not w.status()["live"]) # THE FAILURE THIS EXISTS FOR: the provider accepts the socket and # never delivers — the exact shape of a missing or wrong x-token. # /doctor read `connected` and printed a green tick over a cache # nothing was pushing to, while every trade paid the cold fetch. w._connected = True st = w.status() check("connected with NO slot ever is NOT live", not st["live"]) check("and the reason names the likely cause (auth)", "token" in st["stale_reason"].lower(), st["stale_reason"]) w._last_slot_at = _t.monotonic() st = w.status() check("a slot just now IS live", st["live"]) check("and reports its age", st["secs_since_slot"] is not None) w._last_slot_at = _t.monotonic() - (_STALE_AFTER_SECS + 5) st = w.status() check("a stalled subscription is NOT live", not st["live"]) check("and says the slots stopped, not that auth is wrong", "stopped delivering" in st["stale_reason"], st["stale_reason"]) # /doctor must consume `live`, not `connected`. ch = (Path(__file__).resolve().parent.parent / "modules" / "command_handlers.py").read_text(encoding="utf-8") check('/doctor branches on ws["live"]', 'ws.get("live")' in ch, "reading connected is what produced the green tick over a dead " "subscription") def test_dead_link_detection() -> None: """A half-open connection must become a reconnect, not an outage.""" print("\n\033[1mdead links are detected in seconds, not minutes\033[0m") import modules.ws_state as ws # ping_interval + ping_timeout was 20+20 = up to 40s blind, and the old # 60s reconnect ceiling could add a minute on top: ~100s during which # the socket is open, /doctor says live, and no blockhash arrives. check("ping detection is under 30s total", ws._PING_INTERVAL_SECS + ws._PING_TIMEOUT_SECS <= 30.0, f"{ws._PING_INTERVAL_SECS}+{ws._PING_TIMEOUT_SECS}") check("reconnect backoff is capped low — this is the hot path", ws._RECONNECT_MAX_SECS <= 20.0, str(ws._RECONNECT_MAX_SECS)) # Pings do not catch a peer whose kernel still ACKs while the # application has stopped. Only a read deadline does. check("slot silence forces a reconnect", 0 < ws._SLOT_SILENCE_SECS <= 30.0, str(ws._SLOT_SILENCE_SECS)) src = (Path(__file__).resolve().parent.parent / "modules" / "ws_state.py").read_text(encoding="utf-8") code = "\n".join(l for l in src.splitlines() if not l.lstrip().startswith("#")) check("the read has a deadline, not `async for` forever", "wait_for(" in code and "_SLOT_SILENCE_SECS" in code, "async for waits forever — pings alone miss a half-open socket") # Slots arrive every ~400ms; the silence threshold must be far above # that or a normal gap would flap the connection. check("the threshold is well clear of one slot (~0.4s)", ws._SLOT_SILENCE_SECS >= 5.0, "too tight and a normal jitter becomes a reconnect storm") def main() -> int: print("\n\033[1m══════ ws_state: the pushed cache ══════\033[0m") test_ws_url_override() test_grpc_preflight() test_grpc_blockhash_shape() test_grpc_never_required() test_liveness_not_just_connected() test_dead_link_detection() try: asyncio.run(asyncio.wait_for(run(), timeout=60)) except Exception as exc: # noqa: BLE001 print(f"\n\033[31mharness error: {type(exc).__name__}: {exc}\033[0m") return 1 print() if FAILURES: print(f"\033[31m❌ {len(FAILURES)} of {CHECKS} checks FAILED\033[0m") for f in FAILURES: print(f" • {f}") print() return 1 print(f"\033[32m✅ all {CHECKS} checks passed\033[0m\n") return 0 if __name__ == "__main__": raise SystemExit(main())