Skip to content

Commit

Permalink
Add proper synchronisation to WebSocketTestSession (#2597)
Browse files Browse the repository at this point in the history
* 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 <marcelotryle@gmail.com>
  • Loading branch information
Olocool17 and Kludex authored Jul 20, 2024
1 parent 5f57ef4 commit 357a291
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion starlette/testclient.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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})
Expand Down

0 comments on commit 357a291

Please sign in to comment.