diff --git a/README.md b/README.md index d081b7e..e482467 100644 --- a/README.md +++ b/README.md @@ -194,6 +194,9 @@ except APIError as e: separately with `SUBSCRIPTION_TYPE_ORDER_SNAPSHOT` and a distinct request ID. A successful snapshot ends with an `eof: true` frame; failures use the `error` handler. +For server error frames, both streams emit a `WebSocketError` with optional +`request_id` and `subscription_type` attributes when supplied by the server. + ```python import asyncio import os diff --git a/polymarket_us/errors.py b/polymarket_us/errors.py index 789891d..b969f20 100644 --- a/polymarket_us/errors.py +++ b/polymarket_us/errors.py @@ -123,9 +123,16 @@ class InternalServerError(APIStatusError): class WebSocketError(PolymarketUSError): """WebSocket-related error.""" - def __init__(self, message: str, request_id: str | None = None) -> None: + def __init__( + self, + message: str, + request_id: str | None = None, + *, + subscription_type: str | None = None, + ) -> None: super().__init__(message) self.request_id = request_id + self.subscription_type = subscription_type def __repr__(self) -> str: return f"WebSocketError(message={str(self)!r}, request_id={self.request_id!r})" diff --git a/polymarket_us/websocket/markets.py b/polymarket_us/websocket/markets.py index aac0a05..42d317b 100644 --- a/polymarket_us/websocket/markets.py +++ b/polymarket_us/websocket/markets.py @@ -72,7 +72,11 @@ def _handle_message(self, data: str) -> None: elif "error" in message: self._emit( "error", - WebSocketError(message["error"], message.get("requestId")), + WebSocketError( + message["error"], + message.get("requestId"), + subscription_type=message.get("subscriptionType"), + ), ) elif "marketData" in message: self._emit("market_data", message) diff --git a/polymarket_us/websocket/private.py b/polymarket_us/websocket/private.py index 5efd9e8..e12e699 100644 --- a/polymarket_us/websocket/private.py +++ b/polymarket_us/websocket/private.py @@ -79,7 +79,11 @@ def _handle_message(self, data: str) -> None: elif message.get("error"): self._emit( "error", - WebSocketError(message["error"], message.get("requestId")), + WebSocketError( + message["error"], + message.get("requestId"), + subscription_type=message.get("subscriptionType"), + ), ) elif "orderSubscriptionSnapshot" in message or "ordersSnapshot" in message: self._emit("order_snapshot", message) diff --git a/polymarket_us/websocket/types.py b/polymarket_us/websocket/types.py index 649a798..a531e0b 100644 --- a/polymarket_us/websocket/types.py +++ b/polymarket_us/websocket/types.py @@ -272,6 +272,7 @@ class WebSocketErrorMessage(TypedDict, total=False): requestId: str error: str + subscriptionType: str PrivateMessage = ( diff --git a/pyproject.toml b/pyproject.toml index cbdf8d5..f66243f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "polymarket-us" -version = "2.4.0" +version = "2.4.1" description = "Polymarket US Python SDK" readme = "README.md" license = "MIT" diff --git a/tests/test_errors.py b/tests/test_errors.py index a7ad533..8bc4d4c 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -178,3 +178,16 @@ def test_stores_request_id(self) -> None: """Should store request_id.""" error = WebSocketError("test", request_id="req-123") assert error.request_id == "req-123" + assert error.subscription_type is None + + def test_subscription_context_preserves_existing_error_behavior(self) -> None: + error = WebSocketError( + "subscription failed", "req-123", subscription_type="SUBSCRIPTION_TYPE_ORDER" + ) + assert error.subscription_type == "SUBSCRIPTION_TYPE_ORDER" + assert error.request_id == "req-123" + assert error.args == ("subscription failed",) + assert str(error) == "subscription failed" + assert repr(error) == ( + "WebSocketError(message='subscription failed', request_id='req-123')" + ) diff --git a/tests/test_websocket.py b/tests/test_websocket.py index fa70987..3f81b11 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -2,6 +2,7 @@ import asyncio import base64 +import json from collections.abc import AsyncIterator from typing import Literal @@ -11,9 +12,10 @@ from websockets.http11 import Request from polymarket_us import AuthenticationError, PolymarketUS -from polymarket_us.errors import PolymarketUSError +from polymarket_us.errors import PolymarketUSError, WebSocketError from polymarket_us.websocket import MarketsWebSocket, PrivateWebSocket from polymarket_us.websocket.base import BaseWebSocket +from polymarket_us.websocket.types import WebSocketErrorMessage TEST_SECRET_KEY = "nWGxne/9WmC6hEr0kuwsxERJxWl7MmkZcDusAxyuf2A=" @@ -84,6 +86,50 @@ async def handler(peer: ServerConnection) -> None: finally: await asyncio.wait_for(ws.close(), timeout=5) + @pytest.mark.parametrize( + "subscription_type", + ["known", "SUBSCRIPTION_TYPE_UNSPECIFIED", "SUBSCRIPTION_TYPE_FUTURE", None], + ) + @pytest.mark.parametrize("request_id", ["req-123", None]) + async def test_error_callback_preserves_subscription_context( + self, + connection: tuple[BaseWebSocket, ServerConnection], + stream: Literal["markets", "private"], + subscription_type: str | None, + request_id: str | None, + ) -> None: + ws, peer = connection + if subscription_type == "known": + subscription_type = ( + "SUBSCRIPTION_TYPE_MARKET_DATA" + if stream == "markets" + else "SUBSCRIPTION_TYPE_ORDER" + ) + frame: WebSocketErrorMessage = {"error": "subscription failed"} + if subscription_type is not None: + frame["subscriptionType"] = subscription_type + if request_id is not None: + frame["requestId"] = request_id + errors: list[PolymarketUSError] = [] + messages: list[object] = [] + heartbeat = asyncio.Event() + ws.on("error", errors.append) + ws.on("message", messages.append) + ws.on("heartbeat", heartbeat.set) + + await asyncio.wait_for(peer.send(json.dumps(frame)), timeout=5) + await asyncio.wait_for(peer.send('{"heartbeat": {}}'), timeout=5) + await asyncio.wait_for(heartbeat.wait(), timeout=5) + + assert messages == [frame, {"heartbeat": {}}] + assert len(errors) == 1 + error = errors[0] + assert isinstance(error, WebSocketError) + assert str(error) == "subscription failed" + assert error.request_id == request_id + assert error.subscription_type == subscription_type + assert ws.is_connected + @pytest.mark.parametrize("close_code", [1000, 1001, 1011]) async def test_remote_close_emits_once( self, connection: tuple[BaseWebSocket, ServerConnection], close_code: int diff --git a/tests/types/websocket_errors.py b/tests/types/websocket_errors.py new file mode 100644 index 0000000..8c2de39 --- /dev/null +++ b/tests/types/websocket_errors.py @@ -0,0 +1,17 @@ +"""Typed consumers of WebSocket error context, including unknown subscriptions.""" + +from polymarket_us import WebSocketError +from polymarket_us.websocket import WebSocketErrorMessage + +legacy_frame: WebSocketErrorMessage = {"error": "subscription failed"} +frame: WebSocketErrorMessage = { + "error": "subscription failed", + "requestId": "req-123", + "subscriptionType": "SUBSCRIPTION_TYPE_FUTURE", +} +error = WebSocketError( + frame["error"], + frame.get("requestId"), + subscription_type=frame.get("subscriptionType"), +) +subscription_type: str | None = error.subscription_type