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.
This commit is contained in:
Patryk Zawadzki 2019-04-02 11:52:08 +02:00 committed by Tom Christie
parent 77b84a08c1
commit 3226c7c6f4
2 changed files with 28 additions and 0 deletions

View File

@ -297,6 +297,8 @@ class WebSocketTestSession:
self._send_queue.put(exc) self._send_queue.put(exc)
async def _asgi_receive(self) -> Message: async def _asgi_receive(self) -> Message:
while self._receive_queue.empty():
await asyncio.sleep(0)
return self._receive_queue.get() return self._receive_queue.get()
async def _asgi_send(self, message: Message) -> None: async def _asgi_send(self, message: Message) -> None:

View File

@ -1,8 +1,10 @@
import asyncio
import pytest import pytest
from starlette.applications import Starlette from starlette.applications import Starlette
from starlette.responses import JSONResponse from starlette.responses import JSONResponse
from starlette.testclient import TestClient from starlette.testclient import TestClient
from starlette.websockets import WebSocket, WebSocketDisconnect
mock_service = Starlette() mock_service = Starlette()
@ -86,3 +88,27 @@ def test_testclient_asgi3():
client = TestClient(app) client = TestClient(app)
response = client.get("/") response = client.get("/")
assert response.text == "Hello, world!" 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"}