diff --git a/httpx_ws/_api.py b/httpx_ws/_api.py index 70ef5aa..73d1fa7 100644 --- a/httpx_ws/_api.py +++ b/httpx_ws/_api.py @@ -266,7 +266,7 @@ def receive(self, timeout: float | None = None) -> wsproto.events.Event: A raw [wsproto.events.Event][wsproto.events.Event]. Raises: - queue.Empty: No event was received before the timeout delay. + TimeoutError: No event was received before the timeout delay. WebSocketDisconnect: The server closed the websocket. WebSocketNetworkError: A network error occured. @@ -282,12 +282,15 @@ def receive(self, timeout: float | None = None) -> wsproto.events.Event: try: event = ws.receive(timeout=2.) - except queue.Empty: + except TimeoutError: print("No event received.") except WebSocketDisconnect: print("Connection closed") """ - event = self._events.get(block=True, timeout=timeout) + try: + event = self._events.get(block=True, timeout=timeout) + except queue.Empty as e: + raise TimeoutError from e if isinstance(event, HTTPXWSException): raise event if isinstance(event, wsproto.events.CloseConnection): @@ -307,7 +310,7 @@ def receive_text(self, timeout: float | None = None) -> str: Text data. Raises: - queue.Empty: No event was received before the timeout delay. + TimeoutError: No event was received before the timeout delay. WebSocketDisconnect: The server closed the websocket. WebSocketNetworkError: A network error occured. WebSocketInvalidTypeReceived: The received event was not a text message. @@ -324,7 +327,7 @@ def receive_text(self, timeout: float | None = None) -> str: try: event = ws.receive_text(timeout=2.) - except queue.Empty: + except TimeoutError: print("No text received.") except WebSocketDisconnect: print("Connection closed") @@ -347,7 +350,7 @@ def receive_bytes(self, timeout: float | None = None) -> bytes: Bytes data. Raises: - queue.Empty: No event was received before the timeout delay. + TimeoutError: No event was received before the timeout delay. WebSocketDisconnect: The server closed the websocket. WebSocketNetworkError: A network error occured. WebSocketInvalidTypeReceived: The received event was not a bytes message. @@ -364,7 +367,7 @@ def receive_bytes(self, timeout: float | None = None) -> bytes: try: data = ws.receive_bytes(timeout=2.) - except queue.Empty: + except TimeoutError: print("No data received.") except WebSocketDisconnect: print("Connection closed") @@ -393,7 +396,7 @@ def receive_json( Parsed JSON data. Raises: - queue.Empty: No event was received before the timeout delay. + TimeoutError: No event was received before the timeout delay. WebSocketDisconnect: The server closed the websocket. WebSocketNetworkError: A network error occured. WebSocketInvalidTypeReceived: The received event @@ -411,7 +414,7 @@ def receive_json( try: data = ws.receive_json(timeout=2.) - except queue.Empty: + except TimeoutError: print("No data received.") except WebSocketDisconnect: print("Connection closed") diff --git a/tests/test_api.py b/tests/test_api.py index 230af18..7005f6f 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -258,6 +258,28 @@ def close(self) -> None: with WebSocketSession(stream) as websocket_session: websocket_session.receive() + def test_receive_timeout(self): + class MockNetworkStream(NetworkStream): + def __init__(self) -> None: + self.connection = wsproto.connection.Connection( + wsproto.connection.ConnectionType.SERVER + ) + + def read(self, max_bytes: int, timeout: float | None = None) -> bytes: + time.sleep(0.2) + return b"" + + def write(self, buffer: bytes, timeout: float | None = None) -> None: + pass + + def close(self) -> None: + pass + + stream = MockNetworkStream() + with pytest.raises(TimeoutError): + with WebSocketSession(stream) as websocket_session: + websocket_session.receive(timeout=0.1) + async def test_async_receive_error(self): class AsyncMockNetworkStream(AsyncNetworkStream): def __init__(self) -> None: