Add option timeout to WebSocketTestSession::receive

docs/testclient.md

@@ -156,6 +156,7 @@ May raise `starlette.websockets.WebSocketDisconnect` if the application does not
 * `.receive_bytes()` - Wait for incoming bytestring sent by the application and return it.
 * `.receive_json(mode="text")` - Wait for incoming json data sent by the application and return it. Use `mode="binary"` to receive JSON over binary data frames.
 
+Pass `timeout=timeout_in_seconds` to raise `Empty` when no message is received within `timeout_in_seconds` to prevent blocking.
 May raise `starlette.websockets.WebSocketDisconnect`.
 
 #### Closing the connection

starlette/testclient.py

@@ -156,25 +156,27 @@ def send_json(self, data: typing.Any, mode: str = "text") -> None:
     def close(self, code: int = 1000, reason: typing.Union[str, None] = None) -> None:
         self.send({"type": "websocket.disconnect", "code": code, "reason": reason})
 
-    def receive(self) -> Message:
-        message = self._send_queue.get()
+    def receive(self, timeout: typing.Optional[float] = None) -> Message:
+        message = self._send_queue.get(timeout=timeout)
         if isinstance(message, BaseException):
             raise message
         return message
 
-    def receive_text(self) -> str:
-        message = self.receive()
+    def receive_text(self, timeout: typing.Optional[float] = None) -> str:
+        message = self.receive(timeout=timeout)
         self._raise_on_close(message)
         return typing.cast(str, message["text"])
 
-    def receive_bytes(self) -> bytes:
-        message = self.receive()
+    def receive_bytes(self, timeout: typing.Optional[float] = None) -> bytes:
+        message = self.receive(timeout=timeout)
         self._raise_on_close(message)
         return typing.cast(bytes, message["bytes"])
 
-    def receive_json(self, mode: str = "text") -> typing.Any:
+    def receive_json(
+        self, mode: str = "text", timeout: typing.Optional[float] = None
+    ) -> typing.Any:
         assert mode in ["text", "binary"]
-        message = self.receive()
+        message = self.receive(timeout=timeout)
         self._raise_on_close(message)
         if mode == "text":
             text = message["text"]

tests/test_testclient.py

@@ -2,6 +2,7 @@
 import sys
 from asyncio import current_task as asyncio_current_task
 from contextlib import asynccontextmanager
+from queue import Empty
 from typing import Callable
 
 import anyio
@@ -254,6 +255,20 @@ async def asgi(receive, send):
         assert data == {"message": "test"}
 
 
+def test_websocket_timeout_receive(test_client_factory):
+    def app(scope):
+        async def asgi(receive, send):
+            websocket = WebSocket(scope, receive=receive, send=send)
+            await websocket.accept()
+
+        return asgi
+
+    client = test_client_factory(app)
+    with client.websocket_connect("/") as websocket:
+        with pytest.raises(Empty):
+            websocket.receive_json(timeout=1.0)
+
+
 def test_client(test_client_factory):
     async def app(scope, receive, send):
         client = scope.get("client")