From 87385ada25f644931ed5be46beb16b67dbb93c95 Mon Sep 17 00:00:00 2001 From: harley-poly Date: Wed, 23 Sep 2026 11:31:54 -0400 Subject: [PATCH] fix: emit close after normal WebSocket shutdown --- polymarket_us/websocket/base.py | 2 + pyproject.toml | 2 +- tests/test_websocket.py | 109 ++++++++++++++++++++++++++++++++ 3 files changed, 112 insertions(+), 1 deletion(-) diff --git a/polymarket_us/websocket/base.py b/polymarket_us/websocket/base.py index 001022d..b993350 100644 --- a/polymarket_us/websocket/base.py +++ b/polymarket_us/websocket/base.py @@ -69,6 +69,8 @@ async def _message_loop(self) -> None: self._emit("close") except Exception as e: self._emit("error", PolymarketUSError(str(e))) + else: + self._emit("close") def _handle_message(self, data: str) -> None: """Handle incoming message (override in subclasses).""" diff --git a/pyproject.toml b/pyproject.toml index c129bd8..274deb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "polymarket-us" -version = "1.0.1" +version = "1.0.2" description = "Polymarket US Python SDK" readme = "README.md" license = "MIT" diff --git a/tests/test_websocket.py b/tests/test_websocket.py index 28cb762..fa70987 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -2,6 +2,7 @@ import asyncio import base64 +from collections.abc import AsyncIterator from typing import Literal import pytest @@ -10,7 +11,9 @@ from websockets.http11 import Request from polymarket_us import AuthenticationError, PolymarketUS +from polymarket_us.errors import PolymarketUSError from polymarket_us.websocket import MarketsWebSocket, PrivateWebSocket +from polymarket_us.websocket.base import BaseWebSocket TEST_SECRET_KEY = "nWGxne/9WmC6hEr0kuwsxERJxWl7MmkZcDusAxyuf2A=" @@ -54,6 +57,112 @@ async def handler(connection: ServerConnection) -> None: assert not ws.is_connected +@pytest.mark.parametrize("stream", ["markets", "private"]) +class TestWebSocketLifecycle: + @pytest.fixture + async def connection( + self, stream: Literal["markets", "private"] + ) -> AsyncIterator[tuple[BaseWebSocket, ServerConnection]]: + connections: asyncio.Queue[ServerConnection] = asyncio.Queue() + + async def handler(peer: ServerConnection) -> None: + await connections.put(peer) + await peer.wait_closed() + + async with serve(handler, "127.0.0.1", 0, close_timeout=1) as server: + port = server.sockets[0].getsockname()[1] + with PolymarketUS( + key_id="test-key", + secret_key=TEST_SECRET_KEY, + api_base_url=f"http://127.0.0.1:{port}", + ) as client: + ws = client.ws.markets() if stream == "markets" else client.ws.private() + try: + await asyncio.wait_for(ws.connect(), timeout=5) + peer = await asyncio.wait_for(connections.get(), timeout=5) + yield ws, peer + finally: + await asyncio.wait_for(ws.close(), timeout=5) + + @pytest.mark.parametrize("close_code", [1000, 1001, 1011]) + async def test_remote_close_emits_once( + self, connection: tuple[BaseWebSocket, ServerConnection], close_code: int + ) -> None: + ws, peer = connection + events: list[str] = [] + errors: list[PolymarketUSError] = [] + heartbeat = asyncio.Event() + ws.on("heartbeat", heartbeat.set) + ws.on("close", lambda: events.append("close")) + ws.once("close", lambda: events.append("once")) + ws.on("error", errors.append) + + await asyncio.wait_for(peer.send('{"heartbeat": {}}'), timeout=5) + await asyncio.wait_for(heartbeat.wait(), timeout=5) + await asyncio.wait_for(peer.close(close_code), timeout=5) + assert ws._message_task is not None + await asyncio.wait_for(ws._message_task, timeout=5) + + assert not ws.is_connected + assert events == ["close", "once"] + assert errors == [] + await asyncio.wait_for(ws.close(), timeout=5) + assert events == ["close", "once"] + + async def test_local_close_does_not_emit_close( + self, connection: tuple[BaseWebSocket, ServerConnection] + ) -> None: + ws, peer = connection + events: list[str] = [] + ws.on("close", lambda: events.append("close")) + + await asyncio.wait_for(ws.close(), timeout=5) + await asyncio.wait_for(peer.wait_closed(), timeout=5) + + assert ws._message_task is not None + assert ws._message_task.cancelled() + assert not ws.is_connected + assert events == [] + + async def test_message_task_cancellation_does_not_emit_close( + self, connection: tuple[BaseWebSocket, ServerConnection] + ) -> None: + ws, _ = connection + events: list[str] = [] + ws.on("close", lambda: events.append("close")) + assert ws._message_task is not None + + ws._message_task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(ws._message_task, timeout=5) + + assert ws.is_connected + assert events == [] + + async def test_message_listener_error_does_not_emit_close( + self, connection: tuple[BaseWebSocket, ServerConnection] + ) -> None: + ws, peer = connection + events: list[str] = [] + errors: list[PolymarketUSError] = [] + + def on_heartbeat() -> None: + raise ValueError("listener failed") + + ws.on("heartbeat", on_heartbeat) + ws.on("close", lambda: events.append("close")) + ws.on("error", errors.append) + await asyncio.wait_for(peer.send('{"heartbeat": {}}'), timeout=5) + assert ws._message_task is not None + await asyncio.wait_for(ws._message_task, timeout=5) + + assert ws.is_connected + assert events == [] + assert len(errors) == 1 + assert isinstance(errors[0], PolymarketUSError) + assert str(errors[0]) == "listener failed" + + class TestWebSocketFactory: """Tests for WebSocket factory."""