Cutlery/calibre
1
0
Fork 0
mirror of https://github.com/kovidgoyal/calibre.git synced 2025-07-09 03:04:10 -04:00
Code Issues Packages Projects Releases Wiki Activity

Fix resourcewarnings in websocket tests

Browse Source
This commit is contained in:
Kovid Goyal 2024-08-26 16:08:12 +05:30
parent 60555c1389
commit 08e7cddeb0
No known key found for this signature in database
GPG Key ID: 06BC317B515ACE7C
1 changed files with 41 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

76
src/calibre/srv/tests/web_sockets.py
View File

@ -165,6 +165,12 @@ class WSClient:
reason = reason.encode('utf-8') reason = reason.encode('utf-8')
self.write_frame(1, CLOSE, struct.pack(b'!H', code) + reason) self.write_frame(1, CLOSE, struct.pack(b'!H', code) + reason)
def __enter__(self):
return self
def __exit__(self, *a):
self.socket.close()
class WSTestServer(TestServer): class WSTestServer(TestServer):
@ -184,42 +190,42 @@ class WSTestServer(TestServer):
class WebSocketTest(BaseTest): class WebSocketTest(BaseTest):
def simple_test(self, server, msgs, expected=(), close_code=NORMAL_CLOSE, send_close=True, close_reason=b'NORMAL CLOSE', ignore_send_failures=False): def simple_test(self, server, msgs, expected=(), close_code=NORMAL_CLOSE, send_close=True, close_reason=b'NORMAL CLOSE', ignore_send_failures=False):
client = server.connect() with server.connect() as client:
for msg in msgs: for msg in msgs:
try: try:
if isinstance(msg, dict): if isinstance(msg, dict):
client.write_frame(**msg) client.write_frame(**msg)
else: else:
client.write_message(msg) client.write_message(msg)
except Exception: except Exception:
if not ignore_send_failures: if not ignore_send_failures:
raise raise
expected_messages, expected_controls = [], [] expected_messages, expected_controls = [], []
for ex in expected: for ex in expected:
if isinstance(ex, str): if isinstance(ex, str):
ex = TEXT, ex ex = TEXT, ex
elif isinstance(ex, bytes): elif isinstance(ex, bytes):
ex = BINARY, ex ex = BINARY, ex
elif isinstance(ex, numbers.Integral): elif isinstance(ex, numbers.Integral):
ex = ex, b'' ex = ex, b''
if ex[0] in CONTROL_CODES: if ex[0] in CONTROL_CODES:
expected_controls.append(ex) expected_controls.append(ex)
else: else:
expected_messages.append(ex) expected_messages.append(ex)
if send_close: if send_close:
client.write_close(close_code, close_reason) client.write_close(close_code, close_reason)
try: try:
messages, control_frames = client.read_messages() messages, control_frames = client.read_messages()
except ConnectionAbortedError: except ConnectionAbortedError:
if expected_messages or expected_controls or send_close: if expected_messages or expected_controls or send_close:
raise raise
return return
self.ae(expected_messages, messages) self.ae(expected_messages, messages)
self.assertGreaterEqual(len(control_frames), 1) self.assertGreaterEqual(len(control_frames), 1)
self.ae(expected_controls, control_frames[:-1]) self.ae(expected_controls, control_frames[:-1])
self.ae(control_frames[-1][0], CLOSE) self.ae(control_frames[-1][0], CLOSE)
self.ae(close_code, struct.unpack_from(b'!H', control_frames[-1][1], 0)[0]) self.ae(close_code, struct.unpack_from(b'!H', control_frames[-1][1], 0)[0])
def test_websocket_basic(self): def test_websocket_basic(self):
'Test basic interaction with the websocket server' 'Test basic interaction with the websocket server'