-
-
Notifications
You must be signed in to change notification settings - Fork 954
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Cancel WebSocketTestSession
on close
#2427
Conversation
Kludex
commented
Jan 19, 2024
•
edited
Loading
edited
- Closes Background Thread created by WebSocketTestSession should exit when close() is called #947
4dbfbc7
to
339836e
Compare
starlette/testclient.py
Outdated
async def run_app(tg: anyio.abc.TaskGroup) -> None: | ||
try: | ||
await self.app(self.scope, self._asgi_receive, self._asgi_send) | ||
except anyio.get_cancelled_exc_class(): | ||
... | ||
except BaseException as exc: | ||
self._send_queue.put(exc) | ||
raise | ||
tg.cancel_scope.cancel() | ||
|
||
async with anyio.create_task_group() as tg: | ||
tg.start_soon(run_app, tg) | ||
await self.event.wait() | ||
tg.cancel_scope.cancel() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is the relevant code.
def __exit__(self, *args: typing.Any) -> None: | ||
try: | ||
self.close(1000) | ||
finally: | ||
self.portal.start_task_soon(self._notify_close) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
And this.
tests/test_testclient.py
Outdated
def test_websocket_not_block_on_close(test_client_factory: Callable[..., TestClient]): | ||
def app(scope: Scope): | ||
async def asgi(receive: Receive, send: Send): | ||
websocket = WebSocket(scope, receive=receive, send=send) | ||
await websocket.accept() | ||
while True: | ||
await anyio.sleep(0.1) | ||
|
||
return asgi | ||
|
||
client = test_client_factory(app) | ||
with client.websocket_connect("/") as websocket: | ||
... | ||
assert websocket.event.is_set() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is the relevant test.
@@ -1,16 +1,20 @@ | |||
from __future__ import annotations |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If reviewer complains about this, I can undo the changes.
@encode/maintainers someone? 🙏 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM! This approach with the event flag makes total sense.
starlette/testclient.py
Outdated
except BaseException as exc: | ||
self._send_queue.put(exc) | ||
raise | ||
tg.cancel_scope.cancel() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Shouldn't we also cancel the task group before raising the exception just above?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I did that now.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks.
e7ae532
to
8f78e7a
Compare
* Cancel `WebSocketTestSession` on close * Undo some noise * Fix test * Undo pyproject * Undo anyio bump * Undo changes on test_authentication * Always call cancel scope