From 3226c7c6f4fe8ea8339052696bdc417d44d11524 Mon Sep 17 00:00:00 2001
From: Patryk Zawadzki <81205+patrys@users.noreply.github.com>
Date: Tue, 2 Apr 2019 11:52:08 +0200
Subject: [PATCH] Don't block the event loop in WebSocketTestSession (#459)

It's typical for event-loop-based servers to try to do a blocking
receive in a while loop. Queue.get() is blocking in a synchronous
way and it does not yield control back to the asyncio executor.
Let's explicitly yield control until the queue is no longer empty.
---
 starlette/testclient.py  |  2 ++
 tests/test_testclient.py | 26 ++++++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/starlette/testclient.py b/starlette/testclient.py
index ae46f27b..7679207a 100644
--- a/starlette/testclient.py
+++ b/starlette/testclient.py
@@ -297,6 +297,8 @@ class WebSocketTestSession:
             self._send_queue.put(exc)
 
     async def _asgi_receive(self) -> Message:
+        while self._receive_queue.empty():
+            await asyncio.sleep(0)
         return self._receive_queue.get()
 
     async def _asgi_send(self, message: Message) -> None:
diff --git a/tests/test_testclient.py b/tests/test_testclient.py
index be21c628..1f92c71f 100644
--- a/tests/test_testclient.py
+++ b/tests/test_testclient.py
@@ -1,8 +1,10 @@
+import asyncio
 import pytest
 
 from starlette.applications import Starlette
 from starlette.responses import JSONResponse
 from starlette.testclient import TestClient
+from starlette.websockets import WebSocket, WebSocketDisconnect
 
 mock_service = Starlette()
 
@@ -86,3 +88,27 @@ def test_testclient_asgi3():
     client = TestClient(app)
     response = client.get("/")
     assert response.text == "Hello, world!"
+
+
+def test_websocket_blocking_receive():
+    def app(scope):
+        async def respond(websocket):
+            await websocket.send_json({"message": "test"})
+
+        async def asgi(receive, send):
+            websocket = WebSocket(scope, receive=receive, send=send)
+            await websocket.accept()
+            asyncio.ensure_future(respond(websocket))
+            try:
+                # this will block as the client does not send us data
+                # it should not prevent `respond` from executing though
+                await websocket.receive_json()
+            except WebSocketDisconnect:
+                pass
+
+        return asgi
+
+    client = TestClient(app)
+    with client.websocket_connect("/") as websocket:
+        data = websocket.receive_json()
+        assert data == {"message": "test"}