ResolveSpec/clients/resolvespec-python/tests/test_websocket.py

import asyncio
import json

import pytest
from websockets.asyncio.server import serve

from resolvespec import ResolveSpecError, WebSocketClient


class Server:
    """Minimal in-process WebSocketSpec server."""

    def __init__(self):
        self.received = []
        self.conns = set()
        self.respond = True

    async def handler(self, ws):
        self.conns.add(ws)
        try:
            async for raw in ws:
                msg = json.loads(raw)
                self.received.append(msg)
                if msg["type"] == "ping":
                    await ws.send(json.dumps({"type": "pong"}))
                    continue
                if not self.respond:
                    continue
                await ws.send(json.dumps(self.reply(msg)))
        finally:
            self.conns.discard(ws)

    def reply(self, msg):
        base = {"id": msg["id"], "type": "response", "success": True, "timestamp": "t"}
        if msg["type"] == "subscription" and msg["operation"] == "subscribe":
            return {**base, "data": {"subscription_id": "sub-1"}}
        if msg.get("entity") == "fail":
            return {**base, "success": False, "error": {"code": "bad", "message": "boom"}}
        return {**base, "data": {"echo": msg.get("operation"), "record_id": msg.get("record_id")}}


@pytest.fixture
async def server():
    s = Server()
    async with serve(s.handler, "127.0.0.1", 0) as srv:
        s.url = "ws://127.0.0.1:%d" % srv.sockets[0].getsockname()[1]
        yield s


async def test_operations_and_message_shape(server):
    async with WebSocketClient(server.url, reconnect=False) as c:
        assert c.state == "connected"
        assert await c.read("users", schema="public", record_id="1", limit=5, filters=[{"column": "a", "operator": "eq", "value": 1}]) == {"echo": "read", "record_id": "1"}
        await c.create("users", {"n": 1}, schema="public")
        await c.update("users", "2", {"n": 2})
        await c.delete("users", "3")
        await c.meta("users")
    m = server.received
    assert m[0]["type"] == "request" and m[0]["operation"] == "read"
    assert m[0]["schema"] == "public" and m[0]["record_id"] == "1"
    assert m[0]["options"] == {"filters": [{"column": "a", "operator": "eq", "value": 1}], "limit": 5}
    assert m[1]["data"] == {"n": 1}
    assert m[2]["record_id"] == "2"
    assert [x["operation"] for x in m] == ["read", "create", "update", "delete", "meta"]
    assert "schema" not in m[2]
    assert len({x["id"] for x in m}) == 5


async def test_error_response_raises(server):
    async with WebSocketClient(server.url, reconnect=False) as c:
        with pytest.raises(ResolveSpecError, match="boom") as ei:
            await c.read("fail")
    assert ei.value.code == "bad"


async def test_request_timeout(server):
    server.respond = False
    async with WebSocketClient(server.url, reconnect=False, request_timeout=0.1) as c:
        with pytest.raises(ResolveSpecError, match="timeout"):
            await c.read("users")
        assert not c._pending


async def test_not_connected_raises():
    c = WebSocketClient("ws://127.0.0.1:1")
    with pytest.raises(ResolveSpecError, match="not connected"):
        await c.read("users")


async def test_subscribe_notify_unsubscribe(server):
    got = asyncio.Queue()
    async with WebSocketClient(server.url, reconnect=False) as c:
        sid = await c.subscribe("users", got.put, schema="public", filters=[{"column": "a", "operator": "eq", "value": 1}])
        assert sid == "sub-1"
        assert [s.id for s in c.get_subscriptions()] == ["sub-1"]
        assert server.received[0]["operation"] == "subscribe"
        assert server.received[0]["options"] == {"filters": [{"column": "a", "operator": "eq", "value": 1}]}
        for ws in server.conns:
            await ws.send(json.dumps({"type": "notification", "operation": "create", "subscription_id": "sub-1",
                                      "entity": "users", "data": {"id": 9}, "timestamp": "t"}))
        n = await asyncio.wait_for(got.get(), 2)
        assert n["data"] == {"id": 9}
        await c.unsubscribe("sub-1")
        assert c.get_subscriptions() == []
        assert server.received[-1] == {**server.received[-1], "operation": "unsubscribe", "subscription_id": "sub-1"}


async def test_events_and_heartbeat(server):
    events = []
    c = WebSocketClient(server.url, reconnect=False, heartbeat_interval=0.05)
    c.on("connect", lambda: events.append("connect"))
    c.on("state_change", lambda s: events.append(s))
    c.on("message", lambda m: events.append(("msg", m["type"])))
    await c.connect()
    await asyncio.sleep(0.2)
    await c.close()
    assert events[:3] == ["connecting", "connected", "connect"]
    assert ("msg", "pong") in events
    assert events[-1] == "disconnected"
    assert any(m["type"] == "ping" for m in server.received)
    with pytest.raises(ValueError):
        c.on("bogus", lambda: None)


async def test_reconnect_after_server_drop(server):
    states = []
    c = WebSocketClient(server.url, reconnect=True, reconnect_interval=0.05)
    c.on("state_change", states.append)
    await c.connect()
    for ws in list(server.conns):
        await ws.close()
    for _ in range(100):
        if states.count("connected") >= 2:
            break
        await asyncio.sleep(0.05)
    assert "reconnecting" in states
    assert c.is_connected()
    assert (await c.read("users"))["echo"] == "read"
    await c.close()


async def test_pending_requests_fail_on_disconnect(server):
    server.respond = False
    c = WebSocketClient(server.url, reconnect=False)
    await c.connect()
    task = asyncio.create_task(c.read("users"))
    await asyncio.sleep(0.05)
    for ws in list(server.conns):
        await ws.close()
    with pytest.raises(ResolveSpecError, match="disconnected"):
        await asyncio.wait_for(task, 2)
    await c.close()