From 357a291aa61a813565139f8acafc0b8859f0c571 Mon Sep 17 00:00:00 2001 From: Olocool17 <22843298+Olocool17@users.noreply.github.com> Date: Sat, 20 Jul 2024 09:06:51 +0000 Subject: [PATCH] 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 --- starlette/testclient.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/starlette/testclient.py b/starlette/testclient.py index f17d4e892..bf928d23f 100644 --- a/starlette/testclient.py +++ b/starlette/testclient.py @@ -160,7 +160,8 @@ async def run_app(tg: anyio.abc.TaskGroup) -> None: 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 @@ def _raise_on_close(self, message: Message) -> None: 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})