velvox
/
litecord
mirror of https://gitlab.com/litecord/litecord.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

tests: send correct protocol reply on close event

This commit is contained in:
Luna 2022-01-28 00:28:25 -03:00
parent 94e1f16be5
commit 28299607cd
1 changed files with 9 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tests/test_websocket.py
View File

@ -68,7 +68,7 @@ class AsyncWebsocket:
self.writer.write(self.ws.send(data)) self.writer.write(self.ws.send(data))
await self.writer.drain() await self.writer.drain()
async def recv(self, *, expect=Message, message_str: bool = True): async def recv(self, *, expect=Message, process_event: bool = True):
in_data = await self.reader.read(4096) in_data = await self.reader.read(4096)
if not in_data: if not in_data:
log.info("connection closed (no data)") log.info("connection closed (no data)")
@ -85,6 +85,8 @@ class AsyncWebsocket:
event = next(self.ws.events()) event = next(self.ws.events())
if isinstance(event, CloseConnection): if isinstance(event, CloseConnection):
await self.send(event.response())
if process_event:
raise websockets.ConnectionClosed( raise websockets.ConnectionClosed(
RcvdWrapper(event.code, event.reason), None RcvdWrapper(event.code, event.reason), None
) )
@ -95,7 +97,7 @@ class AsyncWebsocket:
) )
# this keeps compatibility with code written for aaugustin/websockets # this keeps compatibility with code written for aaugustin/websockets
if expect is Message and message_str: if expect is Message and process_event:
return event.data return event.data
return event return event
@ -141,9 +143,9 @@ async def _recv(conn, *, zlib_stream: bool):
zlib_buffer = bytearray() zlib_buffer = bytearray()
while True: while True:
# keep receiving frames until we find the zlib prefix inside # keep receiving frames until we find the zlib prefix inside
# we set message_str to false so that we get the entire event # we set process_event to false so that we get the entire event
# instead of only data # instead of only data
event = await conn.recv(message_str=False) event = await conn.recv(process_event=False)
zlib_buffer.extend(event.data) zlib_buffer.extend(event.data)
if not event.message_finished: if not event.message_finished:
continue continue