scripts/test_ws_state.py

#!/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())