mirror of https://github.com/encode/starlette.git
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:
parent
77b84a08c1
commit
3226c7c6f4
|
@ -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:
|
||||||
|
|
|
@ -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"}
|
||||||
|
|
Loading…
Reference in New Issue