From: Olocool17 <22843298+Olocool17@users.noreply.github.com> Date: Sat, 20 Jul 2024 09:06:51 +0000 (+0000) Subject: Add proper synchronisation to WebSocketTestSession (#2597) X-Git-Tag: 0.38.0~1 X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=357a291aa61a813565139f8acafc0b8859f0c571;p=thirdparty%2Fstarlette.git Add proper synchronisation to WebSocketTestSession (#2597) * Add proper synchronisation to WebSocketTestSession `anyio.sleep(0)` is often used as a way to yield to another task. However, depending on event loop implememtation it is not guaranteed to actually do so in a timely manner. This commit alters this behaviour in _asgi_receive by using `anyio.Event`s as a simple synchronisation primitive, dramatically speeding up the session depending on underlying system/implementation. * Fix mypy type errors Jinja 3.1.4 slightly changed the argument types of FileSystemLoader. * Formatting --------- Co-authored-by: Marcelo Trylesinski --- diff --git a/starlette/testclient.py b/starlette/testclient.py index f17d4e89..bf928d23 100644 --- a/starlette/testclient.py +++ b/starlette/testclient.py @@ -160,7 +160,8 @@ class WebSocketTestSession: async def _asgi_receive(self) -> Message: while self._receive_queue.empty(): - await anyio.sleep(0) + self._queue_event = anyio.Event() + await self._queue_event.wait() return self._receive_queue.get() async def _asgi_send(self, message: Message) -> None: @@ -189,6 +190,8 @@ class WebSocketTestSession: def send(self, message: Message) -> None: self._receive_queue.put(message) + if hasattr(self, "_queue_event"): + self.portal.start_task_soon(self._queue_event.set) def send_text(self, data: str) -> None: self.send({"type": "websocket.receive", "text": data})