Skip to content

Commit

Permalink
Check the disconncet code
Browse files Browse the repository at this point in the history
  • Loading branch information
Kludex committed Nov 1, 2022
1 parent e562633 commit cec0a7d
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions tests/protocols/test_websocket.py
Expand Up @@ -731,15 +731,20 @@ 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):
await self.send({"type": "websocket.accept"})
# 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"))
Expand All @@ -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
Expand Down

0 comments on commit cec0a7d

Please sign in to comment.