diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index 366a4f0bdb..c796e1a696 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -731,7 +731,7 @@ async def test_server_can_read_messages_in_buffer_after_close( ws_protocol_cls, http_protocol_cls ): frames = [] - client_close_connection = asyncio.Event() + disconnect_message = {} class App(WebSocketResponse): async def websocket_connect(self, message): @@ -739,7 +739,12 @@ async def websocket_connect(self, message): # Ensure server doesn't start reading frames from read buffer until # after client has sent close frame, but server is still able to # read these frames - await client_close_connection.wait() + await asyncio.sleep(0.2) + + async def websocket_disconnect(self, message): + nonlocal disconnect_message + print(message) + disconnect_message = message async def websocket_receive(self, message): frames.append(message.get("bytes")) @@ -753,9 +758,9 @@ async def send_text(url): config = Config(app=App, ws=ws_protocol_cls, http=http_protocol_cls, lifespan="off") async with run_server(config): await send_text("ws://127.0.0.1:8000") - client_close_connection.set() assert frames == [b"abc", b"abc", b"abc"] + assert disconnect_message == {"type": "websocket.disconnect", "code": 1000} @pytest.mark.anyio