Raise TimeoutError instead of queue.Empty in WebSocketSession

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")

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: