Initial commit (Clean history)

This commit is contained in:
anhduy-tech
2025-12-30 11:27:14 +07:00
commit ef48c93de0
19255 changed files with 3248867 additions and 0 deletions

View File

@@ -0,0 +1,6 @@
HTTP/1.1 101 WebSocket Protocol Handshake
Connection: Upgrade
Upgrade: WebSocket
Sec-WebSocket-Accept: Kxep+hNu9n51529fGidYu7a3wO0=
some_header: something

View File

@@ -0,0 +1,6 @@
HTTP/1.1 101 WebSocket Protocol Handshake
Connection: Upgrade
Upgrade WebSocket
Sec-WebSocket-Accept: Kxep+hNu9n51529fGidYu7a3wO0=
some_header: something

View File

@@ -0,0 +1,8 @@
HTTP/1.1 101 WebSocket Protocol Handshake
Connection: Upgrade, Keep-Alive
Upgrade: WebSocket
Sec-WebSocket-Accept: Kxep+hNu9n51529fGidYu7a3wO0=
Set-Cookie: Token=ABCDE
Set-Cookie: Token=FGHIJ
some_header: something

View File

@@ -0,0 +1,23 @@
#!/usr/bin/env python
# From https://github.com/aaugustin/websockets/blob/main/example/echo.py
import asyncio
import os
import websockets
LOCAL_WS_SERVER_PORT = int(os.environ.get("LOCAL_WS_SERVER_PORT", "8765"))
async def echo(websocket):
async for message in websocket:
await websocket.send(message)
async def main():
async with websockets.serve(echo, "localhost", LOCAL_WS_SERVER_PORT):
await asyncio.Future() # run forever
asyncio.run(main())

View File

@@ -0,0 +1,125 @@
# -*- coding: utf-8 -*-
#
import unittest
from websocket._abnf import ABNF, frame_buffer
from websocket._exceptions import WebSocketProtocolException
"""
test_abnf.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class ABNFTest(unittest.TestCase):
def test_init(self):
a = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_PING)
self.assertEqual(a.fin, 0)
self.assertEqual(a.rsv1, 0)
self.assertEqual(a.rsv2, 0)
self.assertEqual(a.rsv3, 0)
self.assertEqual(a.opcode, 9)
self.assertEqual(a.data, "")
a_bad = ABNF(0, 1, 0, 0, opcode=77)
self.assertEqual(a_bad.rsv1, 1)
self.assertEqual(a_bad.opcode, 77)
def test_validate(self):
a_invalid_ping = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_PING)
self.assertRaises(
WebSocketProtocolException,
a_invalid_ping.validate,
skip_utf8_validation=False,
)
a_bad_rsv_value = ABNF(0, 1, 0, 0, opcode=ABNF.OPCODE_TEXT)
self.assertRaises(
WebSocketProtocolException,
a_bad_rsv_value.validate,
skip_utf8_validation=False,
)
a_bad_opcode = ABNF(0, 0, 0, 0, opcode=77)
self.assertRaises(
WebSocketProtocolException,
a_bad_opcode.validate,
skip_utf8_validation=False,
)
a_bad_close_frame = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_CLOSE, data=b"\x01")
self.assertRaises(
WebSocketProtocolException,
a_bad_close_frame.validate,
skip_utf8_validation=False,
)
a_bad_close_frame_2 = ABNF(
0, 0, 0, 0, opcode=ABNF.OPCODE_CLOSE, data=b"\x01\x8a\xaa\xff\xdd"
)
self.assertRaises(
WebSocketProtocolException,
a_bad_close_frame_2.validate,
skip_utf8_validation=False,
)
a_bad_close_frame_3 = ABNF(
0, 0, 0, 0, opcode=ABNF.OPCODE_CLOSE, data=b"\x03\xe7"
)
self.assertRaises(
WebSocketProtocolException,
a_bad_close_frame_3.validate,
skip_utf8_validation=True,
)
def test_mask(self):
abnf_none_data = ABNF(
0, 0, 0, 0, opcode=ABNF.OPCODE_PING, mask_value=1, data=None
)
bytes_val = b"aaaa"
self.assertEqual(abnf_none_data._get_masked(bytes_val), bytes_val)
abnf_str_data = ABNF(
0, 0, 0, 0, opcode=ABNF.OPCODE_PING, mask_value=1, data="a"
)
self.assertEqual(abnf_str_data._get_masked(bytes_val), b"aaaa\x00")
def test_format(self):
abnf_bad_rsv_bits = ABNF(2, 0, 0, 0, opcode=ABNF.OPCODE_TEXT)
self.assertRaises(ValueError, abnf_bad_rsv_bits.format)
abnf_bad_opcode = ABNF(0, 0, 0, 0, opcode=5)
self.assertRaises(ValueError, abnf_bad_opcode.format)
abnf_length_10 = ABNF(0, 0, 0, 0, opcode=ABNF.OPCODE_TEXT, data="abcdefghij")
self.assertEqual(b"\x01", abnf_length_10.format()[0].to_bytes(1, "big"))
self.assertEqual(b"\x8a", abnf_length_10.format()[1].to_bytes(1, "big"))
self.assertEqual("fin=0 opcode=1 data=abcdefghij", abnf_length_10.__str__())
abnf_length_20 = ABNF(
0, 0, 0, 0, opcode=ABNF.OPCODE_BINARY, data="abcdefghijabcdefghij"
)
self.assertEqual(b"\x02", abnf_length_20.format()[0].to_bytes(1, "big"))
self.assertEqual(b"\x94", abnf_length_20.format()[1].to_bytes(1, "big"))
abnf_no_mask = ABNF(
0, 0, 0, 0, opcode=ABNF.OPCODE_TEXT, mask_value=0, data=b"\x01\x8a\xcc"
)
self.assertEqual(b"\x01\x03\x01\x8a\xcc", abnf_no_mask.format())
def test_frame_buffer(self):
fb = frame_buffer(0, True)
self.assertEqual(fb.recv, 0)
self.assertEqual(fb.skip_utf8_validation, True)
fb.clear
self.assertEqual(fb.header, None)
self.assertEqual(fb.length, None)
self.assertEqual(fb.mask_value, None)
self.assertEqual(fb.has_mask(), False)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,395 @@
# -*- coding: utf-8 -*-
#
import os
import os.path
import ssl
import threading
import unittest
import websocket as ws
"""
test_app.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
# Skip test to access the internet unless TEST_WITH_INTERNET == 1
TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"
# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -1
LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
TRACEABLE = True
class WebSocketAppTest(unittest.TestCase):
class NotSetYet:
"""A marker class for signalling that a value hasn't been set yet."""
def setUp(self):
ws.enableTrace(TRACEABLE)
WebSocketAppTest.keep_running_open = WebSocketAppTest.NotSetYet()
WebSocketAppTest.keep_running_close = WebSocketAppTest.NotSetYet()
WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet()
WebSocketAppTest.on_error_data = WebSocketAppTest.NotSetYet()
def tearDown(self):
WebSocketAppTest.keep_running_open = WebSocketAppTest.NotSetYet()
WebSocketAppTest.keep_running_close = WebSocketAppTest.NotSetYet()
WebSocketAppTest.get_mask_key_id = WebSocketAppTest.NotSetYet()
WebSocketAppTest.on_error_data = WebSocketAppTest.NotSetYet()
def close(self):
pass
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_keep_running(self):
"""A WebSocketApp should keep running as long as its self.keep_running
is not False (in the boolean context).
"""
def on_open(self, *args, **kwargs):
"""Set the keep_running flag for later inspection and immediately
close the connection.
"""
self.send("hello!")
WebSocketAppTest.keep_running_open = self.keep_running
self.keep_running = False
def on_message(_, message):
print(message)
self.close()
def on_close(self, *args, **kwargs):
"""Set the keep_running flag for the test to use."""
WebSocketAppTest.keep_running_close = self.keep_running
app = ws.WebSocketApp(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
on_open=on_open,
on_close=on_close,
on_message=on_message,
)
app.run_forever()
# @unittest.skipUnless(TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled")
@unittest.skipUnless(False, "Test disabled for now (requires rel)")
def test_run_forever_dispatcher(self):
"""A WebSocketApp should keep running as long as its self.keep_running
is not False (in the boolean context).
"""
def on_open(self, *args, **kwargs):
"""Send a message, receive, and send one more"""
self.send("hello!")
self.recv()
self.send("goodbye!")
def on_message(_, message):
print(message)
self.close()
app = ws.WebSocketApp(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
on_open=on_open,
on_message=on_message,
)
app.run_forever(dispatcher="Dispatcher") # doesn't work
# app.run_forever(dispatcher=rel) # would work
# rel.dispatch()
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_run_forever_teardown_clean_exit(self):
"""The WebSocketApp.run_forever() method should return `False` when the application ends gracefully."""
app = ws.WebSocketApp(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
threading.Timer(interval=0.2, function=app.close).start()
teardown = app.run_forever()
self.assertEqual(teardown, False)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_sock_mask_key(self):
"""A WebSocketApp should forward the received mask_key function down
to the actual socket.
"""
def my_mask_key_func():
return "\x00\x00\x00\x00"
app = ws.WebSocketApp(
"wss://api-pub.bitfinex.com/ws/1", get_mask_key=my_mask_key_func
)
# if numpy is installed, this assertion fail
# Note: We can't use 'is' for comparing the functions directly, need to use 'id'.
self.assertEqual(id(app.get_mask_key), id(my_mask_key_func))
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_invalid_ping_interval_ping_timeout(self):
"""Test exception handling if ping_interval < ping_timeout"""
def on_ping(app, _):
print("Got a ping!")
app.close()
def on_pong(app, _):
print("Got a pong! No need to respond")
app.close()
app = ws.WebSocketApp(
"wss://api-pub.bitfinex.com/ws/1", on_ping=on_ping, on_pong=on_pong
)
self.assertRaises(
ws.WebSocketException,
app.run_forever,
ping_interval=1,
ping_timeout=2,
sslopt={"cert_reqs": ssl.CERT_NONE},
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_ping_interval(self):
"""Test WebSocketApp proper ping functionality"""
def on_ping(app, _):
print("Got a ping!")
app.close()
def on_pong(app, _):
print("Got a pong! No need to respond")
app.close()
app = ws.WebSocketApp(
"wss://api-pub.bitfinex.com/ws/1", on_ping=on_ping, on_pong=on_pong
)
app.run_forever(
ping_interval=2, ping_timeout=1, sslopt={"cert_reqs": ssl.CERT_NONE}
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_opcode_close(self):
"""Test WebSocketApp close opcode"""
app = ws.WebSocketApp("wss://tsock.us1.twilio.com/v3/wsconnect")
app.run_forever(ping_interval=2, ping_timeout=1, ping_payload="Ping payload")
# This is commented out because the URL no longer responds in the expected way
# @unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
# def testOpcodeBinary(self):
# """ Test WebSocketApp binary opcode
# """
# app = ws.WebSocketApp('wss://streaming.vn.teslamotors.com/streaming/')
# app.run_forever(ping_interval=2, ping_timeout=1, ping_payload="Ping payload")
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_bad_ping_interval(self):
"""A WebSocketApp handling of negative ping_interval"""
app = ws.WebSocketApp("wss://api-pub.bitfinex.com/ws/1")
self.assertRaises(
ws.WebSocketException,
app.run_forever,
ping_interval=-5,
sslopt={"cert_reqs": ssl.CERT_NONE},
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_bad_ping_timeout(self):
"""A WebSocketApp handling of negative ping_timeout"""
app = ws.WebSocketApp("wss://api-pub.bitfinex.com/ws/1")
self.assertRaises(
ws.WebSocketException,
app.run_forever,
ping_timeout=-3,
sslopt={"cert_reqs": ssl.CERT_NONE},
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_close_status_code(self):
"""Test extraction of close frame status code and close reason in WebSocketApp"""
def on_close(wsapp, close_status_code, close_msg):
print("on_close reached")
app = ws.WebSocketApp(
"wss://tsock.us1.twilio.com/v3/wsconnect", on_close=on_close
)
closeframe = ws.ABNF(
opcode=ws.ABNF.OPCODE_CLOSE, data=b"\x03\xe8no-init-from-client"
)
self.assertEqual([1000, "no-init-from-client"], app._get_close_args(closeframe))
closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b"")
self.assertEqual([None, None], app._get_close_args(closeframe))
app2 = ws.WebSocketApp("wss://tsock.us1.twilio.com/v3/wsconnect")
closeframe = ws.ABNF(opcode=ws.ABNF.OPCODE_CLOSE, data=b"")
self.assertEqual([None, None], app2._get_close_args(closeframe))
self.assertRaises(
ws.WebSocketConnectionClosedException,
app.send,
data="test if connection is closed",
)
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_callback_function_exception(self):
"""Test callback function exception handling"""
exc = None
passed_app = None
def on_open(app):
raise RuntimeError("Callback failed")
def on_error(app, err):
nonlocal passed_app
passed_app = app
nonlocal exc
exc = err
def on_pong(app, _):
app.close()
app = ws.WebSocketApp(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
on_open=on_open,
on_error=on_error,
on_pong=on_pong,
)
app.run_forever(ping_interval=2, ping_timeout=1)
self.assertEqual(passed_app, app)
self.assertIsInstance(exc, RuntimeError)
self.assertEqual(str(exc), "Callback failed")
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_callback_method_exception(self):
"""Test callback method exception handling"""
class Callbacks:
def __init__(self):
self.exc = None
self.passed_app = None
self.app = ws.WebSocketApp(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
on_open=self.on_open,
on_error=self.on_error,
on_pong=self.on_pong,
)
self.app.run_forever(ping_interval=2, ping_timeout=1)
def on_open(self, _):
raise RuntimeError("Callback failed")
def on_error(self, app, err):
self.passed_app = app
self.exc = err
def on_pong(self, app, _):
app.close()
callbacks = Callbacks()
self.assertEqual(callbacks.passed_app, callbacks.app)
self.assertIsInstance(callbacks.exc, RuntimeError)
self.assertEqual(str(callbacks.exc), "Callback failed")
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_reconnect(self):
"""Test reconnect"""
pong_count = 0
exc = None
def on_error(_, err):
nonlocal exc
exc = err
def on_pong(app, _):
nonlocal pong_count
pong_count += 1
if pong_count == 1:
# First pong, shutdown socket, enforce read error
app.sock.shutdown()
if pong_count >= 2:
# Got second pong after reconnect
app.close()
app = ws.WebSocketApp(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}", on_pong=on_pong, on_error=on_error
)
app.run_forever(ping_interval=2, ping_timeout=1, reconnect=3)
self.assertEqual(pong_count, 2)
self.assertIsInstance(exc, ws.WebSocketTimeoutException)
self.assertEqual(str(exc), "ping/pong timed out")
def test_dispatcher_selection_default(self):
"""Test default dispatcher selection"""
app = ws.WebSocketApp("ws://example.com")
# Test default dispatcher (non-SSL)
dispatcher = app.create_dispatcher(ping_timeout=10, is_ssl=False)
self.assertIsInstance(dispatcher, ws._dispatcher.Dispatcher)
def test_dispatcher_selection_ssl(self):
"""Test SSL dispatcher selection"""
app = ws.WebSocketApp("wss://example.com")
# Test SSL dispatcher
dispatcher = app.create_dispatcher(ping_timeout=10, is_ssl=True)
self.assertIsInstance(dispatcher, ws._dispatcher.SSLDispatcher)
def test_dispatcher_selection_custom(self):
"""Test custom dispatcher selection"""
from unittest.mock import Mock
app = ws.WebSocketApp("ws://example.com")
custom_dispatcher = Mock()
handle_disconnect = Mock()
# Test wrapped dispatcher with custom dispatcher
dispatcher = app.create_dispatcher(
ping_timeout=10,
dispatcher=custom_dispatcher,
handleDisconnect=handle_disconnect,
)
self.assertIsInstance(dispatcher, ws._dispatcher.WrappedDispatcher)
self.assertEqual(dispatcher.dispatcher, custom_dispatcher)
self.assertEqual(dispatcher.handleDisconnect, handle_disconnect)
def test_dispatcher_selection_no_ping_timeout(self):
"""Test dispatcher selection without ping timeout"""
app = ws.WebSocketApp("ws://example.com")
# Test with None ping_timeout (should default to 10)
dispatcher = app.create_dispatcher(ping_timeout=None, is_ssl=False)
self.assertIsInstance(dispatcher, ws._dispatcher.Dispatcher)
self.assertEqual(dispatcher.ping_timeout, 10)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,123 @@
import unittest
from websocket._cookiejar import SimpleCookieJar
"""
test_cookiejar.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class CookieJarTest(unittest.TestCase):
def test_add(self):
cookie_jar = SimpleCookieJar()
cookie_jar.add("")
self.assertFalse(
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
)
cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b")
self.assertFalse(
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
)
cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b; domain=.abc")
self.assertTrue(".abc" in cookie_jar.jar)
cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b; domain=abc")
self.assertTrue(".abc" in cookie_jar.jar)
self.assertTrue("abc" not in cookie_jar.jar)
cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b; c=d; domain=abc")
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
self.assertEqual(cookie_jar.get(None), "")
cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b; c=d; domain=abc")
cookie_jar.add("e=f; domain=abc")
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d; e=f")
cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b; c=d; domain=abc")
cookie_jar.add("e=f; domain=.abc")
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d; e=f")
cookie_jar = SimpleCookieJar()
cookie_jar.add("a=b; c=d; domain=abc")
cookie_jar.add("e=f; domain=xyz")
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
self.assertEqual(cookie_jar.get("xyz"), "e=f")
self.assertEqual(cookie_jar.get("something"), "")
def test_set(self):
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b")
self.assertFalse(
cookie_jar.jar, "Cookie with no domain should not be added to the jar"
)
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; domain=.abc")
self.assertTrue(".abc" in cookie_jar.jar)
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; domain=abc")
self.assertTrue(".abc" in cookie_jar.jar)
self.assertTrue("abc" not in cookie_jar.jar)
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; c=d; domain=abc")
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; c=d; domain=abc")
cookie_jar.set("e=f; domain=abc")
self.assertEqual(cookie_jar.get("abc"), "e=f")
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; c=d; domain=abc")
cookie_jar.set("e=f; domain=.abc")
self.assertEqual(cookie_jar.get("abc"), "e=f")
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; c=d; domain=abc")
cookie_jar.set("e=f; domain=xyz")
self.assertEqual(cookie_jar.get("abc"), "a=b; c=d")
self.assertEqual(cookie_jar.get("xyz"), "e=f")
self.assertEqual(cookie_jar.get("something"), "")
def test_get(self):
cookie_jar = SimpleCookieJar()
cookie_jar.set("a=b; c=d; domain=abc.com")
self.assertEqual(cookie_jar.get("abc.com"), "a=b; c=d")
self.assertEqual(cookie_jar.get("x.abc.com"), "a=b; c=d")
self.assertEqual(cookie_jar.get("abc.com.es"), "")
self.assertEqual(cookie_jar.get("xabc.com"), "")
cookie_jar.set("a=b; c=d; domain=.abc.com")
self.assertEqual(cookie_jar.get("abc.com"), "a=b; c=d")
self.assertEqual(cookie_jar.get("x.abc.com"), "a=b; c=d")
self.assertEqual(cookie_jar.get("abc.com.es"), "")
self.assertEqual(cookie_jar.get("xabc.com"), "")
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,385 @@
# -*- coding: utf-8 -*-
import socket
import unittest
from unittest.mock import Mock, patch, MagicMock
import threading
import time
import websocket
from websocket._dispatcher import (
Dispatcher,
DispatcherBase,
SSLDispatcher,
WrappedDispatcher,
)
"""
test_dispatcher.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class MockApp:
"""Mock WebSocketApp for testing"""
def __init__(self):
self.keep_running = True
self.sock = Mock()
self.sock.sock = Mock()
class MockSocket:
"""Mock socket for testing"""
def __init__(self):
self.pending_return = False
def pending(self):
return self.pending_return
class MockDispatcher:
"""Mock external dispatcher for WrappedDispatcher testing"""
def __init__(self):
self.signal_calls = []
self.abort_calls = []
self.read_calls = []
self.buffwrite_calls = []
self.timeout_calls = []
def signal(self, sig, handler):
self.signal_calls.append((sig, handler))
def abort(self):
self.abort_calls.append(True)
def read(self, sock, callback):
self.read_calls.append((sock, callback))
def buffwrite(self, sock, data, send_func, disconnect_handler):
self.buffwrite_calls.append((sock, data, send_func, disconnect_handler))
def timeout(self, seconds, callback, *args):
self.timeout_calls.append((seconds, callback, args))
class DispatcherTest(unittest.TestCase):
def setUp(self):
self.app = MockApp()
def test_dispatcher_base_init(self):
"""Test DispatcherBase initialization"""
dispatcher = DispatcherBase(self.app, 30.0)
self.assertEqual(dispatcher.app, self.app)
self.assertEqual(dispatcher.ping_timeout, 30.0)
def test_dispatcher_base_timeout(self):
"""Test DispatcherBase timeout method"""
dispatcher = DispatcherBase(self.app, 30.0)
callback = Mock()
# Test with seconds=None (should call callback immediately)
dispatcher.timeout(None, callback)
callback.assert_called_once()
# Test with seconds > 0 (would sleep in real implementation)
callback.reset_mock()
start_time = time.time()
dispatcher.timeout(0.1, callback)
elapsed = time.time() - start_time
callback.assert_called_once()
self.assertGreaterEqual(elapsed, 0.05) # Allow some tolerance
def test_dispatcher_base_reconnect(self):
"""Test DispatcherBase reconnect method"""
dispatcher = DispatcherBase(self.app, 30.0)
reconnector = Mock()
# Test normal reconnect
dispatcher.reconnect(1, reconnector)
reconnector.assert_called_once_with(reconnecting=True)
# Test reconnect with KeyboardInterrupt
reconnector.reset_mock()
reconnector.side_effect = KeyboardInterrupt("User interrupted")
with self.assertRaises(KeyboardInterrupt):
dispatcher.reconnect(1, reconnector)
def test_dispatcher_base_send(self):
"""Test DispatcherBase send method"""
dispatcher = DispatcherBase(self.app, 30.0)
mock_sock = Mock()
test_data = b"test data"
with patch("websocket._dispatcher.send") as mock_send:
mock_send.return_value = len(test_data)
result = dispatcher.send(mock_sock, test_data)
mock_send.assert_called_once_with(mock_sock, test_data)
self.assertEqual(result, len(test_data))
def test_dispatcher_read(self):
"""Test Dispatcher read method"""
dispatcher = Dispatcher(self.app, 5.0)
read_callback = Mock(return_value=True)
check_callback = Mock()
mock_sock = Mock()
# Mock the selector to control the loop
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
# Make select return immediately (timeout)
mock_selector.select.return_value = []
# Stop after first iteration
def side_effect(*args):
self.app.keep_running = False
return []
mock_selector.select.side_effect = side_effect
dispatcher.read(mock_sock, read_callback, check_callback)
# Verify selector was used correctly
mock_selector.register.assert_called()
mock_selector.select.assert_called_with(5.0)
mock_selector.close.assert_called()
check_callback.assert_called()
def test_dispatcher_read_with_data(self):
"""Test Dispatcher read method when data is available"""
dispatcher = Dispatcher(self.app, 5.0)
read_callback = Mock(return_value=True)
check_callback = Mock()
mock_sock = Mock()
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
# First call returns data, second call stops the loop
call_count = 0
def select_side_effect(*args):
nonlocal call_count
call_count += 1
if call_count == 1:
return [True] # Data available
else:
self.app.keep_running = False
return []
mock_selector.select.side_effect = select_side_effect
dispatcher.read(mock_sock, read_callback, check_callback)
read_callback.assert_called()
check_callback.assert_called()
def test_ssl_dispatcher_read(self):
"""Test SSLDispatcher read method"""
dispatcher = SSLDispatcher(self.app, 5.0)
read_callback = Mock(return_value=True)
check_callback = Mock()
# Mock socket with pending data
mock_ssl_sock = MockSocket()
self.app.sock.sock = mock_ssl_sock
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = []
# Stop after first iteration
def side_effect(*args):
self.app.keep_running = False
return []
mock_selector.select.side_effect = side_effect
dispatcher.read(None, read_callback, check_callback)
mock_selector.register.assert_called()
check_callback.assert_called()
def test_ssl_dispatcher_select_with_pending(self):
"""Test SSLDispatcher select method with pending data"""
dispatcher = SSLDispatcher(self.app, 5.0)
mock_ssl_sock = MockSocket()
mock_ssl_sock.pending_return = True
self.app.sock.sock = mock_ssl_sock
mock_selector = Mock()
result = dispatcher.select(None, mock_selector)
# When pending() returns True, should return [sock]
self.assertEqual(result, [mock_ssl_sock])
def test_ssl_dispatcher_select_without_pending(self):
"""Test SSLDispatcher select method without pending data"""
dispatcher = SSLDispatcher(self.app, 5.0)
mock_ssl_sock = MockSocket()
mock_ssl_sock.pending_return = False
self.app.sock.sock = mock_ssl_sock
mock_selector = Mock()
mock_selector.select.return_value = [(mock_ssl_sock, None)]
result = dispatcher.select(None, mock_selector)
# Should return the first element of first result tuple
self.assertEqual(result, mock_ssl_sock)
mock_selector.select.assert_called_with(5.0)
def test_ssl_dispatcher_select_no_results(self):
"""Test SSLDispatcher select method with no results"""
dispatcher = SSLDispatcher(self.app, 5.0)
mock_ssl_sock = MockSocket()
mock_ssl_sock.pending_return = False
self.app.sock.sock = mock_ssl_sock
mock_selector = Mock()
mock_selector.select.return_value = []
result = dispatcher.select(None, mock_selector)
# Should return None when no results (function doesn't return anything when len(r) == 0)
self.assertIsNone(result)
def test_wrapped_dispatcher_init(self):
"""Test WrappedDispatcher initialization"""
mock_dispatcher = MockDispatcher()
handle_disconnect = Mock()
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
self.assertEqual(wrapped.app, self.app)
self.assertEqual(wrapped.ping_timeout, 10.0)
self.assertEqual(wrapped.dispatcher, mock_dispatcher)
self.assertEqual(wrapped.handleDisconnect, handle_disconnect)
# Should have set up signal handler
self.assertEqual(len(mock_dispatcher.signal_calls), 1)
sig, handler = mock_dispatcher.signal_calls[0]
self.assertEqual(sig, 2) # SIGINT
self.assertEqual(handler, mock_dispatcher.abort)
def test_wrapped_dispatcher_read(self):
"""Test WrappedDispatcher read method"""
mock_dispatcher = MockDispatcher()
handle_disconnect = Mock()
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
mock_sock = Mock()
read_callback = Mock()
check_callback = Mock()
wrapped.read(mock_sock, read_callback, check_callback)
# Should delegate to wrapped dispatcher
self.assertEqual(len(mock_dispatcher.read_calls), 1)
self.assertEqual(mock_dispatcher.read_calls[0], (mock_sock, read_callback))
# Should call timeout for ping_timeout
self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
timeout_call = mock_dispatcher.timeout_calls[0]
self.assertEqual(timeout_call[0], 10.0) # timeout seconds
self.assertEqual(timeout_call[1], check_callback) # callback
def test_wrapped_dispatcher_read_no_ping_timeout(self):
"""Test WrappedDispatcher read method without ping timeout"""
mock_dispatcher = MockDispatcher()
handle_disconnect = Mock()
wrapped = WrappedDispatcher(self.app, None, mock_dispatcher, handle_disconnect)
mock_sock = Mock()
read_callback = Mock()
check_callback = Mock()
wrapped.read(mock_sock, read_callback, check_callback)
# Should delegate to wrapped dispatcher
self.assertEqual(len(mock_dispatcher.read_calls), 1)
# Should NOT call timeout when ping_timeout is None
self.assertEqual(len(mock_dispatcher.timeout_calls), 0)
def test_wrapped_dispatcher_send(self):
"""Test WrappedDispatcher send method"""
mock_dispatcher = MockDispatcher()
handle_disconnect = Mock()
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
mock_sock = Mock()
test_data = b"test data"
with patch("websocket._dispatcher.send") as mock_send:
result = wrapped.send(mock_sock, test_data)
# Should delegate to dispatcher.buffwrite
self.assertEqual(len(mock_dispatcher.buffwrite_calls), 1)
call = mock_dispatcher.buffwrite_calls[0]
self.assertEqual(call[0], mock_sock)
self.assertEqual(call[1], test_data)
self.assertEqual(call[2], mock_send)
self.assertEqual(call[3], handle_disconnect)
# Should return data length
self.assertEqual(result, len(test_data))
def test_wrapped_dispatcher_timeout(self):
"""Test WrappedDispatcher timeout method"""
mock_dispatcher = MockDispatcher()
handle_disconnect = Mock()
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
callback = Mock()
args = ("arg1", "arg2")
wrapped.timeout(5.0, callback, *args)
# Should delegate to wrapped dispatcher
self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
call = mock_dispatcher.timeout_calls[0]
self.assertEqual(call[0], 5.0)
self.assertEqual(call[1], callback)
self.assertEqual(call[2], args)
def test_wrapped_dispatcher_reconnect(self):
"""Test WrappedDispatcher reconnect method"""
mock_dispatcher = MockDispatcher()
handle_disconnect = Mock()
wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
reconnector = Mock()
wrapped.reconnect(3, reconnector)
# Should delegate to timeout method with reconnect=True
self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
call = mock_dispatcher.timeout_calls[0]
self.assertEqual(call[0], 3)
self.assertEqual(call[1], reconnector)
self.assertEqual(call[2], (True,))
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,158 @@
# -*- coding: utf-8 -*-
import unittest
from unittest.mock import Mock, patch
from websocket._handshake import _get_resp_headers
from websocket._exceptions import WebSocketBadStatusException
from websocket._ssl_compat import SSLError
"""
test_handshake_large_response.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class HandshakeLargeResponseTest(unittest.TestCase):
def test_large_error_response_chunked_reading(self):
"""Test that large HTTP error responses during handshake are read in chunks"""
# Mock socket
mock_sock = Mock()
# Create a large error response body (> 16KB)
large_response = b"Error details: " + b"A" * 20000 # 20KB+ response
# Track recv calls to ensure chunking
recv_calls = []
def mock_recv(sock, bufsize):
recv_calls.append(bufsize)
# Simulate SSL error if trying to read > 16KB at once
if bufsize > 16384:
raise SSLError("[SSL: BAD_LENGTH] unknown error")
return large_response[:bufsize]
# Mock read_headers to return error status with large content-length
with patch("websocket._handshake.read_headers") as mock_read_headers:
mock_read_headers.return_value = (
400, # Bad request status
{"content-length": str(len(large_response))},
"Bad Request",
)
# Mock the recv function to track calls
with patch("websocket._socket.recv", side_effect=mock_recv):
# This should not raise SSLError, but should raise WebSocketBadStatusException
with self.assertRaises(WebSocketBadStatusException) as cm:
_get_resp_headers(mock_sock)
# Verify the response body was included in the exception
self.assertIn(
b"Error details:",
(
cm.exception.args[0].encode()
if isinstance(cm.exception.args[0], str)
else cm.exception.args[0]
),
)
# Verify chunked reading was used (multiple recv calls, none > 16KB)
self.assertGreater(len(recv_calls), 1)
self.assertTrue(all(call <= 16384 for call in recv_calls))
def test_handshake_ssl_large_response_protection(self):
"""Test that the fix prevents SSL BAD_LENGTH errors during handshake"""
mock_sock = Mock()
# Large content that would trigger SSL error if read all at once
large_content = b"X" * 32768 # 32KB
chunks_returned = 0
def mock_recv_chunked(sock, bufsize):
nonlocal chunks_returned
# Return data in chunks, simulating successful chunked reading
chunk_start = chunks_returned * 16384
chunk_end = min(chunk_start + bufsize, len(large_content))
result = large_content[chunk_start:chunk_end]
chunks_returned += 1 if result else 0
return result
with patch("websocket._handshake.read_headers") as mock_read_headers:
mock_read_headers.return_value = (
500, # Server error
{"content-length": str(len(large_content))},
"Internal Server Error",
)
with patch("websocket._socket.recv", side_effect=mock_recv_chunked):
# Should handle large response without SSL errors
with self.assertRaises(WebSocketBadStatusException) as cm:
_get_resp_headers(mock_sock)
# Verify the complete response was captured
exception_str = str(cm.exception)
# Response body should be in the exception message
self.assertIn("XXXXX", exception_str) # Part of the large content
def test_handshake_normal_small_response(self):
"""Test that normal small responses still work correctly"""
mock_sock = Mock()
small_response = b"Small error message"
def mock_recv(sock, bufsize):
return small_response
with patch("websocket._handshake.read_headers") as mock_read_headers:
mock_read_headers.return_value = (
404, # Not found
{"content-length": str(len(small_response))},
"Not Found",
)
with patch("websocket._socket.recv", side_effect=mock_recv):
with self.assertRaises(WebSocketBadStatusException) as cm:
_get_resp_headers(mock_sock)
# Verify small response is handled correctly
self.assertIn("Small error message", str(cm.exception))
def test_handshake_no_content_length(self):
"""Test handshake error response without content-length header"""
mock_sock = Mock()
with patch("websocket._handshake.read_headers") as mock_read_headers:
mock_read_headers.return_value = (
403, # Forbidden
{}, # No content-length header
"Forbidden",
)
# Should raise exception without trying to read response body
with self.assertRaises(WebSocketBadStatusException) as cm:
_get_resp_headers(mock_sock)
# Should mention status but not have response body
exception_str = str(cm.exception)
self.assertIn("403", exception_str)
self.assertIn("Forbidden", exception_str)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,370 @@
# -*- coding: utf-8 -*-
#
import os
import os.path
import socket
import ssl
import unittest
import websocket
from websocket._exceptions import WebSocketProxyException, WebSocketException
from websocket._http import (
_get_addrinfo_list,
_start_proxied_socket,
_tunnel,
connect,
proxy_info,
read_headers,
HAVE_PYTHON_SOCKS,
)
"""
test_http.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
try:
from python_socks._errors import ProxyConnectionError, ProxyError, ProxyTimeoutError
except:
from websocket._http import ProxyConnectionError, ProxyError, ProxyTimeoutError
# Skip test to access the internet unless TEST_WITH_INTERNET == 1
TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"
TEST_WITH_PROXY = os.environ.get("TEST_WITH_PROXY", "0") == "1"
# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -1
LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
class SockMock:
def __init__(self):
self.data = []
self.sent = []
def add_packet(self, data):
self.data.append(data)
def gettimeout(self):
return None
def recv(self, bufsize):
if self.data:
e = self.data.pop(0)
if isinstance(e, Exception):
raise e
if len(e) > bufsize:
self.data.insert(0, e[bufsize:])
return e[:bufsize]
def send(self, data):
self.sent.append(data)
return len(data)
def close(self):
pass
class HeaderSockMock(SockMock):
def __init__(self, fname):
SockMock.__init__(self)
path = os.path.join(os.path.dirname(__file__), fname)
with open(path, "rb") as f:
self.add_packet(f.read())
class OptsList:
def __init__(self):
self.timeout = 1
self.sockopt = []
self.sslopt = {"cert_reqs": ssl.CERT_NONE}
class HttpTest(unittest.TestCase):
def test_read_header(self):
status, header, _ = read_headers(HeaderSockMock("data/header01.txt"))
self.assertEqual(status, 101)
self.assertEqual(header["connection"], "Upgrade")
# header02.txt is intentionally malformed
self.assertRaises(
WebSocketException, read_headers, HeaderSockMock("data/header02.txt")
)
def test_tunnel(self):
self.assertRaises(
WebSocketProxyException,
_tunnel,
HeaderSockMock("data/header01.txt"),
"example.com",
80,
("username", "password"),
)
self.assertRaises(
WebSocketProxyException,
_tunnel,
HeaderSockMock("data/header02.txt"),
"example.com",
80,
("username", "password"),
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_connect(self):
# Not currently testing an actual proxy connection, so just check whether proxy errors are raised. This requires internet for a DNS lookup
if HAVE_PYTHON_SOCKS:
# Need this check, otherwise case where python_socks is not installed triggers
# websocket._exceptions.WebSocketException: Python Socks is needed for SOCKS proxying but is not available
self.assertRaises(
(ProxyTimeoutError, OSError),
_start_proxied_socket,
"wss://example.com",
OptsList(),
proxy_info(
http_proxy_host="example.com",
http_proxy_port="8080",
proxy_type="socks4",
http_proxy_timeout=1,
),
)
self.assertRaises(
(ProxyTimeoutError, OSError),
_start_proxied_socket,
"wss://example.com",
OptsList(),
proxy_info(
http_proxy_host="example.com",
http_proxy_port="8080",
proxy_type="socks4a",
http_proxy_timeout=1,
),
)
self.assertRaises(
(ProxyTimeoutError, OSError),
_start_proxied_socket,
"wss://example.com",
OptsList(),
proxy_info(
http_proxy_host="example.com",
http_proxy_port="8080",
proxy_type="socks5",
http_proxy_timeout=1,
),
)
self.assertRaises(
(ProxyTimeoutError, OSError),
_start_proxied_socket,
"wss://example.com",
OptsList(),
proxy_info(
http_proxy_host="example.com",
http_proxy_port="8080",
proxy_type="socks5h",
http_proxy_timeout=1,
),
)
self.assertRaises(
ProxyConnectionError,
connect,
"wss://example.com",
OptsList(),
proxy_info(
http_proxy_host="127.0.0.1",
http_proxy_port=9999,
proxy_type="socks4",
http_proxy_timeout=1,
),
None,
)
self.assertRaises(
TypeError,
_get_addrinfo_list,
None,
80,
True,
proxy_info(
http_proxy_host="127.0.0.1", http_proxy_port="9999", proxy_type="http"
),
)
self.assertRaises(
TypeError,
_get_addrinfo_list,
None,
80,
True,
proxy_info(
http_proxy_host="127.0.0.1", http_proxy_port="9999", proxy_type="http"
),
)
self.assertRaises(
socket.timeout,
connect,
"wss://google.com",
OptsList(),
proxy_info(
http_proxy_host="8.8.8.8",
http_proxy_port=9999,
proxy_type="http",
http_proxy_timeout=1,
),
None,
)
self.assertEqual(
connect(
"wss://google.com",
OptsList(),
proxy_info(
http_proxy_host="8.8.8.8", http_proxy_port=8080, proxy_type="http"
),
True,
),
(True, ("google.com", 443, "/")),
)
# The following test fails on Mac OS with a gaierror, not an OverflowError
# self.assertRaises(OverflowError, connect, "wss://example.com", OptsList(), proxy_info(http_proxy_host="127.0.0.1", http_proxy_port=99999, proxy_type="socks4", timeout=2), False)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
@unittest.skipUnless(
TEST_WITH_PROXY, "This test requires a HTTP proxy to be running on port 8899"
)
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_proxy_connect(self):
ws = websocket.WebSocket()
ws.connect(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
http_proxy_host="127.0.0.1",
http_proxy_port="8899",
proxy_type="http",
)
ws.send("Hello, Server")
server_response = ws.recv()
self.assertEqual(server_response, "Hello, Server")
# self.assertEqual(_start_proxied_socket("wss://api.bitfinex.com/ws/2", OptsList(), proxy_info(http_proxy_host="127.0.0.1", http_proxy_port="8899", proxy_type="http"))[1], ("api.bitfinex.com", 443, '/ws/2'))
self.assertEqual(
_get_addrinfo_list(
"api.bitfinex.com",
443,
True,
proxy_info(
http_proxy_host="127.0.0.1",
http_proxy_port="8899",
proxy_type="http",
),
),
(
socket.getaddrinfo(
"127.0.0.1", 8899, 0, socket.SOCK_STREAM, socket.SOL_TCP
),
True,
None,
),
)
self.assertEqual(
connect(
"wss://api.bitfinex.com/ws/2",
OptsList(),
proxy_info(
http_proxy_host="127.0.0.1", http_proxy_port=8899, proxy_type="http"
),
None,
)[1],
("api.bitfinex.com", 443, "/ws/2"),
)
# TODO: Test SOCKS4 and SOCK5 proxies with unit tests
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_sslopt(self):
ssloptions = {
"check_hostname": False,
"server_hostname": "ServerName",
"ssl_version": ssl.PROTOCOL_TLS_CLIENT,
"ciphers": "TLS_AES_256_GCM_SHA384:TLS_CHACHA20_POLY1305_SHA256:\
TLS_AES_128_GCM_SHA256:ECDHE-ECDSA-AES256-GCM-SHA384:\
ECDHE-RSA-AES256-GCM-SHA384:DHE-RSA-AES256-GCM-SHA384:\
ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-CHACHA20-POLY1305:\
DHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-AES128-GCM-SHA256:\
ECDHE-RSA-AES128-GCM-SHA256:DHE-RSA-AES128-GCM-SHA256:\
ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA384:\
DHE-RSA-AES256-SHA256:ECDHE-ECDSA-AES128-SHA256:\
ECDHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA256:\
ECDHE-ECDSA-AES256-SHA:ECDHE-RSA-AES256-SHA",
"ecdh_curve": "prime256v1",
}
ws_ssl1 = websocket.WebSocket(sslopt=ssloptions)
ws_ssl1.connect("wss://api.bitfinex.com/ws/2")
ws_ssl1.send("Hello")
ws_ssl1.close()
ws_ssl2 = websocket.WebSocket(sslopt={"check_hostname": True})
ws_ssl2.connect("wss://api.bitfinex.com/ws/2")
ws_ssl2.close
def test_proxy_info(self):
self.assertEqual(
proxy_info(
http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"
).proxy_protocol,
"http",
)
self.assertRaises(
ProxyError,
proxy_info,
http_proxy_host="127.0.0.1",
http_proxy_port="8080",
proxy_type="badval",
)
self.assertEqual(
proxy_info(
http_proxy_host="example.com", http_proxy_port="8080", proxy_type="http"
).proxy_host,
"example.com",
)
self.assertEqual(
proxy_info(
http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"
).proxy_port,
"8080",
)
self.assertEqual(
proxy_info(
http_proxy_host="127.0.0.1", http_proxy_port="8080", proxy_type="http"
).auth,
None,
)
self.assertEqual(
proxy_info(
http_proxy_host="127.0.0.1",
http_proxy_port="8080",
proxy_type="http",
http_proxy_auth=("my_username123", "my_pass321"),
).auth[0],
"my_username123",
)
self.assertEqual(
proxy_info(
http_proxy_host="127.0.0.1",
http_proxy_port="8080",
proxy_type="http",
http_proxy_auth=("my_username123", "my_pass321"),
).auth[1],
"my_pass321",
)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,273 @@
# -*- coding: utf-8 -*-
import unittest
import struct
from unittest.mock import Mock, patch, MagicMock
from websocket._abnf import ABNF
from websocket._core import WebSocket
from websocket._exceptions import WebSocketProtocolException, WebSocketPayloadException
from websocket._ssl_compat import SSLError
"""
test_large_payloads.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class LargePayloadTest(unittest.TestCase):
def test_frame_length_encoding_boundaries(self):
"""Test WebSocket frame length encoding at various boundaries"""
# Test length encoding boundaries as per RFC 6455
test_cases = [
(125, "Single byte length"), # Max for 7-bit length
(126, "Two byte length start"), # Start of 16-bit length
(127, "Two byte length"),
(65535, "Two byte length max"), # Max for 16-bit length
(65536, "Eight byte length start"), # Start of 64-bit length
(16384, "16KB boundary"), # The problematic size
(16385, "Just over 16KB"),
(32768, "32KB"),
(131072, "128KB"),
]
for length, description in test_cases:
with self.subTest(length=length, description=description):
# Create payload of specified length
payload = b"A" * length
# Create frame
frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)
# Verify frame can be formatted without error
formatted = frame.format()
# Verify the frame header is correctly structured
self.assertIsInstance(formatted, bytes)
self.assertTrue(len(formatted) >= length) # Header + payload
# Verify payload length is preserved
self.assertEqual(len(frame.data), length)
def test_recv_large_payload_chunked(self):
"""Test receiving large payloads in chunks (simulating the 16KB recv issue)"""
# Create a large payload that would trigger chunked reading
large_payload = b"B" * 32768 # 32KB
# Mock recv function that returns data in 16KB chunks
chunks = []
chunk_size = 16384
for i in range(0, len(large_payload), chunk_size):
chunks.append(large_payload[i : i + chunk_size])
call_count = 0
def mock_recv(bufsize):
nonlocal call_count
if call_count >= len(chunks):
return b""
result = chunks[call_count]
call_count += 1
return result
# Test the frame buffer's recv_strict method
from websocket._abnf import frame_buffer
fb = frame_buffer(mock_recv, skip_utf8_validation=True)
# This should handle large payloads by chunking
result = fb.recv_strict(len(large_payload))
self.assertEqual(result, large_payload)
# Verify multiple recv calls were made
self.assertGreater(call_count, 1)
def test_ssl_large_payload_simulation(self):
"""Simulate SSL BAD_LENGTH error scenario"""
# This test demonstrates that the 16KB limit in frame buffer protects against SSL issues
payload_size = 16385
recv_calls = []
def mock_recv_with_ssl_limit(bufsize):
recv_calls.append(bufsize)
# This simulates the SSL issue: BAD_LENGTH when trying to recv > 16KB
if bufsize > 16384:
raise SSLError("[SSL: BAD_LENGTH] unknown error")
return b"C" * min(bufsize, 16384)
from websocket._abnf import frame_buffer
fb = frame_buffer(mock_recv_with_ssl_limit, skip_utf8_validation=True)
# The frame buffer handles this correctly by chunking recv calls
result = fb.recv_strict(payload_size)
# Verify it worked and chunked the calls properly
self.assertEqual(len(result), payload_size)
# Verify no single recv call was > 16KB
self.assertTrue(all(call <= 16384 for call in recv_calls))
# Verify multiple calls were made
self.assertGreater(len(recv_calls), 1)
def test_frame_format_large_payloads(self):
"""Test frame formatting with various large payload sizes"""
# Test sizes around potential problem areas
test_sizes = [16383, 16384, 16385, 32768, 65535, 65536]
for size in test_sizes:
with self.subTest(size=size):
payload = b"D" * size
frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)
# Should not raise any exceptions
formatted = frame.format()
# Verify structure
self.assertIsInstance(formatted, bytes)
self.assertEqual(len(frame.data), size)
# Verify length encoding is correct based on size
# Note: frames from create_frame() include masking by default (4 extra bytes)
mask_size = 4 # WebSocket frames are masked by default
if size < ABNF.LENGTH_7: # < 126
# Length should be encoded in single byte
expected_header_size = (
2 + mask_size
) # 1 byte opcode + 1 byte length + 4 byte mask
elif size < ABNF.LENGTH_16: # < 65536
# Length should be encoded in 2 bytes
expected_header_size = (
4 + mask_size
) # 1 byte opcode + 1 byte marker + 2 bytes length + 4 byte mask
else:
# Length should be encoded in 8 bytes
expected_header_size = (
10 + mask_size
) # 1 byte opcode + 1 byte marker + 8 bytes length + 4 byte mask
self.assertEqual(len(formatted), expected_header_size + size)
def test_send_large_payload_chunking(self):
"""Test that large payloads are sent in chunks to avoid SSL issues"""
mock_sock = Mock()
# Track how data is sent
sent_chunks = []
def mock_send(data):
sent_chunks.append(len(data))
return len(data)
mock_sock.send = mock_send
mock_sock.gettimeout.return_value = 30.0
# Create WebSocket with mocked socket
ws = WebSocket()
ws.sock = mock_sock
ws.connected = True
# Create large payload
large_payload = b"E" * 32768 # 32KB
# Send the payload
with patch("websocket._core.send") as mock_send_func:
mock_send_func.side_effect = lambda sock, data: len(data)
# This should work without SSL errors
result = ws.send_binary(large_payload)
# Verify payload was accepted
self.assertGreater(result, 0)
def test_utf8_validation_large_text(self):
"""Test UTF-8 validation with large text payloads"""
# Create large valid UTF-8 text
large_text = "Hello 世界! " * 2000 # About 26KB with Unicode
# Test frame creation
frame = ABNF.create_frame(large_text, ABNF.OPCODE_TEXT)
# Should not raise validation errors
formatted = frame.format()
self.assertIsInstance(formatted, bytes)
# Test with close frame that has invalid UTF-8 (this is what validate() actually checks)
invalid_utf8_close_data = struct.pack("!H", 1000) + b"\xff\xfe invalid utf8"
# Create close frame with invalid UTF-8 data
frame = ABNF(1, 0, 0, 0, ABNF.OPCODE_CLOSE, 1, invalid_utf8_close_data)
# Validation should catch the invalid UTF-8 in close frame reason
with self.assertRaises(WebSocketProtocolException):
frame.validate(skip_utf8_validation=False)
def test_frame_buffer_edge_cases(self):
"""Test frame buffer with edge cases that could trigger bugs"""
# Test scenario: exactly 16KB payload split across recv calls
payload_16k = b"F" * 16384
# Simulate receiving in smaller chunks
chunks = [payload_16k[i : i + 4096] for i in range(0, len(payload_16k), 4096)]
call_count = 0
def mock_recv(bufsize):
nonlocal call_count
if call_count >= len(chunks):
return b""
result = chunks[call_count]
call_count += 1
return result
from websocket._abnf import frame_buffer
fb = frame_buffer(mock_recv, skip_utf8_validation=True)
result = fb.recv_strict(16384)
self.assertEqual(result, payload_16k)
# Verify multiple recv calls were made
self.assertEqual(call_count, 4) # 16KB / 4KB = 4 chunks
def test_max_frame_size_limits(self):
"""Test behavior at WebSocket maximum frame size limits"""
# Test just under the maximum theoretical frame size
# (This is a very large test, so we'll use a smaller representative size)
# Test with a reasonably large payload that represents the issue
large_size = 1024 * 1024 # 1MB
payload = b"G" * large_size
# This should work without issues
frame = ABNF.create_frame(payload, ABNF.OPCODE_BINARY)
# Verify the frame can be formatted
formatted = frame.format()
self.assertIsInstance(formatted, bytes)
# Verify payload is preserved
self.assertEqual(len(frame.data), large_size)
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,357 @@
# -*- coding: utf-8 -*-
import errno
import socket
import unittest
from unittest.mock import Mock, patch, MagicMock
import time
from websocket._socket import recv, recv_line, send, DEFAULT_SOCKET_OPTION
from websocket._ssl_compat import (
SSLError,
SSLEOFError,
SSLWantWriteError,
SSLWantReadError,
)
from websocket._exceptions import (
WebSocketTimeoutException,
WebSocketConnectionClosedException,
)
"""
test_socket.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class SocketTest(unittest.TestCase):
def test_default_socket_option(self):
"""Test DEFAULT_SOCKET_OPTION contains expected options"""
self.assertIsInstance(DEFAULT_SOCKET_OPTION, list)
self.assertGreater(len(DEFAULT_SOCKET_OPTION), 0)
# Should contain TCP_NODELAY option
tcp_nodelay_found = any(
opt[1] == socket.TCP_NODELAY for opt in DEFAULT_SOCKET_OPTION
)
self.assertTrue(tcp_nodelay_found)
def test_recv_normal(self):
"""Test normal recv operation"""
mock_sock = Mock()
mock_sock.recv.return_value = b"test data"
result = recv(mock_sock, 9)
self.assertEqual(result, b"test data")
mock_sock.recv.assert_called_once_with(9)
def test_recv_timeout_error(self):
"""Test recv with TimeoutError"""
mock_sock = Mock()
mock_sock.recv.side_effect = TimeoutError("Connection timed out")
with self.assertRaises(WebSocketTimeoutException) as cm:
recv(mock_sock, 9)
self.assertEqual(str(cm.exception), "Connection timed out")
def test_recv_socket_timeout(self):
"""Test recv with socket.timeout"""
mock_sock = Mock()
timeout_exc = socket.timeout("Socket timed out")
timeout_exc.args = ("Socket timed out",)
mock_sock.recv.side_effect = timeout_exc
mock_sock.gettimeout.return_value = 30.0
with self.assertRaises(WebSocketTimeoutException) as cm:
recv(mock_sock, 9)
# In Python 3.10+, socket.timeout is a subclass of TimeoutError
# so it's caught by the TimeoutError handler with hardcoded message
# In Python 3.9, socket.timeout is caught by socket.timeout handler
# which preserves the original message
import sys
if sys.version_info >= (3, 10):
self.assertEqual(str(cm.exception), "Connection timed out")
else:
self.assertEqual(str(cm.exception), "Socket timed out")
def test_recv_ssl_timeout(self):
"""Test recv with SSL timeout error"""
mock_sock = Mock()
ssl_exc = SSLError("The operation timed out")
ssl_exc.args = ("The operation timed out",)
mock_sock.recv.side_effect = ssl_exc
with self.assertRaises(WebSocketTimeoutException) as cm:
recv(mock_sock, 9)
self.assertEqual(str(cm.exception), "The operation timed out")
def test_recv_ssl_non_timeout_error(self):
"""Test recv with SSL non-timeout error"""
mock_sock = Mock()
ssl_exc = SSLError("SSL certificate error")
ssl_exc.args = ("SSL certificate error",)
mock_sock.recv.side_effect = ssl_exc
# Should re-raise the original SSL error
with self.assertRaises(SSLError):
recv(mock_sock, 9)
def test_recv_empty_response(self):
"""Test recv with empty response (connection closed)"""
mock_sock = Mock()
mock_sock.recv.return_value = b""
with self.assertRaises(WebSocketConnectionClosedException) as cm:
recv(mock_sock, 9)
self.assertEqual(str(cm.exception), "Connection to remote host was lost.")
def test_recv_ssl_want_read_error(self):
"""Test recv with SSLWantReadError (should retry)"""
mock_sock = Mock()
# First call raises SSLWantReadError, second call succeeds
mock_sock.recv.side_effect = [SSLWantReadError(), b"data after retry"]
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Ready to read
result = recv(mock_sock, 100)
self.assertEqual(result, b"data after retry")
mock_selector.register.assert_called()
mock_selector.close.assert_called()
def test_recv_ssl_want_read_timeout(self):
"""Test recv with SSLWantReadError that times out"""
mock_sock = Mock()
mock_sock.recv.side_effect = SSLWantReadError()
mock_sock.gettimeout.return_value = 1.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [] # Timeout
with self.assertRaises(WebSocketTimeoutException):
recv(mock_sock, 100)
def test_recv_line(self):
"""Test recv_line functionality"""
mock_sock = Mock()
# Mock recv to return one character at a time
recv_calls = [b"H", b"e", b"l", b"l", b"o", b"\n"]
with patch("websocket._socket.recv", side_effect=recv_calls) as mock_recv:
result = recv_line(mock_sock)
self.assertEqual(result, b"Hello\n")
self.assertEqual(mock_recv.call_count, 6)
def test_send_normal(self):
"""Test normal send operation"""
mock_sock = Mock()
mock_sock.send.return_value = 9
mock_sock.gettimeout.return_value = 30.0
result = send(mock_sock, b"test data")
self.assertEqual(result, 9)
mock_sock.send.assert_called_with(b"test data")
def test_send_zero_timeout(self):
"""Test send with zero timeout (non-blocking)"""
mock_sock = Mock()
mock_sock.send.return_value = 9
mock_sock.gettimeout.return_value = 0
result = send(mock_sock, b"test data")
self.assertEqual(result, 9)
mock_sock.send.assert_called_once_with(b"test data")
def test_send_ssl_eof_error(self):
"""Test send with SSLEOFError"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
mock_sock.send.side_effect = SSLEOFError("Connection closed")
with self.assertRaises(WebSocketConnectionClosedException) as cm:
send(mock_sock, b"test data")
self.assertEqual(str(cm.exception), "socket is already closed.")
def test_send_ssl_want_write_error(self):
"""Test send with SSLWantWriteError (should retry)"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
# First call raises SSLWantWriteError, second call succeeds
mock_sock.send.side_effect = [SSLWantWriteError(), 9]
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Ready to write
result = send(mock_sock, b"test data")
self.assertEqual(result, 9)
mock_selector.register.assert_called()
mock_selector.close.assert_called()
def test_send_socket_eagain_error(self):
"""Test send with EAGAIN error (should retry)"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
# Create socket error with EAGAIN
eagain_error = socket.error("Resource temporarily unavailable")
eagain_error.errno = errno.EAGAIN
eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")
# First call raises EAGAIN, second call succeeds
mock_sock.send.side_effect = [eagain_error, 9]
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Ready to write
result = send(mock_sock, b"test data")
self.assertEqual(result, 9)
def test_send_socket_ewouldblock_error(self):
"""Test send with EWOULDBLOCK error (should retry)"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
# Create socket error with EWOULDBLOCK
ewouldblock_error = socket.error("Operation would block")
ewouldblock_error.errno = errno.EWOULDBLOCK
ewouldblock_error.args = (errno.EWOULDBLOCK, "Operation would block")
# First call raises EWOULDBLOCK, second call succeeds
mock_sock.send.side_effect = [ewouldblock_error, 9]
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Ready to write
result = send(mock_sock, b"test data")
self.assertEqual(result, 9)
def test_send_socket_other_error(self):
"""Test send with other socket error (should raise)"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
# Create socket error with different errno
other_error = socket.error("Connection reset by peer")
other_error.errno = errno.ECONNRESET
other_error.args = (errno.ECONNRESET, "Connection reset by peer")
mock_sock.send.side_effect = other_error
with self.assertRaises(socket.error):
send(mock_sock, b"test data")
def test_send_socket_error_no_errno(self):
"""Test send with socket error that has no errno"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
# Create socket error without errno attribute
no_errno_error = socket.error("Generic socket error")
no_errno_error.args = ("Generic socket error",)
mock_sock.send.side_effect = no_errno_error
with self.assertRaises(socket.error):
send(mock_sock, b"test data")
def test_send_write_timeout(self):
"""Test send write operation timeout"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
# First call raises EAGAIN
eagain_error = socket.error("Resource temporarily unavailable")
eagain_error.errno = errno.EAGAIN
eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")
mock_sock.send.side_effect = eagain_error
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [] # Timeout - nothing ready
result = send(mock_sock, b"test data")
# Should return 0 when write times out
self.assertEqual(result, 0)
def test_send_string_data(self):
"""Test send with string data (should be encoded)"""
mock_sock = Mock()
mock_sock.send.return_value = 9
mock_sock.gettimeout.return_value = 30.0
result = send(mock_sock, "test data")
self.assertEqual(result, 9)
mock_sock.send.assert_called_with(b"test data")
def test_send_partial_send_retry(self):
"""Test send retry mechanism"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 30.0
# Create a scenario where send succeeds after selector retry
eagain_error = socket.error("Resource temporarily unavailable")
eagain_error.errno = errno.EAGAIN
eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")
# Mock the internal _send function behavior
mock_sock.send.side_effect = [eagain_error, 9]
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Socket ready for writing
result = send(mock_sock, b"test data")
self.assertEqual(result, 9)
# Verify selector was used for retry mechanism
mock_selector.register.assert_called()
mock_selector.select.assert_called()
mock_selector.close.assert_called()
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,160 @@
# -*- coding: utf-8 -*-
import errno
import socket
import unittest
from unittest.mock import Mock, patch
from websocket._socket import recv
from websocket._ssl_compat import SSLWantReadError
from websocket._exceptions import (
WebSocketTimeoutException,
WebSocketConnectionClosedException,
)
"""
test_socket_bugs.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class SocketBugsTest(unittest.TestCase):
"""Test bugs found in socket handling logic"""
def test_bug_implicit_none_return_from_ssl_want_read_fixed(self):
"""
BUG #5 FIX VERIFICATION: Test SSLWantReadError timeout now raises correct exception
Bug was in _socket.py:100-101 - SSLWantReadError except block returned None implicitly
Fixed: Now properly handles timeout with WebSocketTimeoutException
"""
mock_sock = Mock()
mock_sock.recv.side_effect = SSLWantReadError()
mock_sock.gettimeout.return_value = 1.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [] # Timeout - no data ready
with self.assertRaises(WebSocketTimeoutException) as cm:
recv(mock_sock, 100)
# Verify correct timeout exception and message
self.assertIn("Connection timed out waiting for data", str(cm.exception))
def test_bug_implicit_none_return_from_socket_error_fixed(self):
"""
BUG #5 FIX VERIFICATION: Test that socket.error with EAGAIN now handles timeout correctly
Bug was in _socket.py:102-105 - socket.error except block returned None implicitly
Fixed: Now properly handles timeout with WebSocketTimeoutException
"""
mock_sock = Mock()
# Create socket error with EAGAIN (should be retried)
eagain_error = OSError(errno.EAGAIN, "Resource temporarily unavailable")
# First call raises EAGAIN, selector times out on retry
mock_sock.recv.side_effect = eagain_error
mock_sock.gettimeout.return_value = 1.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [] # Timeout - no data ready
with self.assertRaises(WebSocketTimeoutException) as cm:
recv(mock_sock, 100)
# Verify correct timeout exception and message
self.assertIn("Connection timed out waiting for data", str(cm.exception))
def test_bug_wrong_exception_for_selector_timeout_fixed(self):
"""
BUG #6 FIX VERIFICATION: Test that selector timeout now raises correct exception type
Bug was in _socket.py:115 returning None for timeout, treated as connection error
Fixed: Now raises WebSocketTimeoutException directly
"""
mock_sock = Mock()
mock_sock.recv.side_effect = SSLWantReadError() # Trigger retry path
mock_sock.gettimeout.return_value = 1.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [] # TIMEOUT - this is key!
with self.assertRaises(WebSocketTimeoutException) as cm:
recv(mock_sock, 100)
# Verify it's the correct timeout exception with proper message
self.assertIn("Connection timed out waiting for data", str(cm.exception))
# This proves the fix works:
# 1. selector.select() returns [] (timeout)
# 2. _recv() now raises WebSocketTimeoutException directly
# 3. No more misclassification as connection closed error!
def test_socket_timeout_exception_handling(self):
"""
Test that socket.timeout exceptions are properly handled
"""
mock_sock = Mock()
mock_sock.gettimeout.return_value = 1.0
# Simulate a real socket.timeout scenario
mock_sock.recv.side_effect = socket.timeout("Operation timed out")
# This works correctly - socket.timeout raises WebSocketTimeoutException
with self.assertRaises(WebSocketTimeoutException) as cm:
recv(mock_sock, 100)
# In Python 3.10+, socket.timeout is a subclass of TimeoutError
# so it's caught by the TimeoutError handler with hardcoded message
# In Python 3.9, socket.timeout is caught by socket.timeout handler
# which preserves the original message
import sys
if sys.version_info >= (3, 10):
self.assertIn("Connection timed out", str(cm.exception))
else:
self.assertIn("Operation timed out", str(cm.exception))
def test_correct_ssl_want_read_retry_behavior(self):
"""Test the correct behavior when SSLWantReadError is properly handled"""
mock_sock = Mock()
# First call raises SSLWantReadError, second call succeeds
mock_sock.recv.side_effect = [SSLWantReadError(), b"data after retry"]
mock_sock.gettimeout.return_value = 1.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Data ready after wait
# This should work correctly
result = recv(mock_sock, 100)
self.assertEqual(result, b"data after retry")
# Selector should be used for retry
mock_selector.register.assert_called()
mock_selector.select.assert_called()
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,91 @@
# -*- coding: utf-8 -*-
import sys
import unittest
from unittest.mock import patch
"""
test_ssl_compat.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class SSLCompatTest(unittest.TestCase):
def test_ssl_available(self):
"""Test that SSL is available in normal conditions"""
import websocket._ssl_compat as ssl_compat
# In normal conditions, SSL should be available
self.assertTrue(ssl_compat.HAVE_SSL)
self.assertIsNotNone(ssl_compat.ssl)
# SSL exception classes should be available
self.assertTrue(hasattr(ssl_compat, "SSLError"))
self.assertTrue(hasattr(ssl_compat, "SSLEOFError"))
self.assertTrue(hasattr(ssl_compat, "SSLWantReadError"))
self.assertTrue(hasattr(ssl_compat, "SSLWantWriteError"))
def test_ssl_not_available(self):
"""Test fallback behavior when SSL is not available"""
# Remove ssl_compat from modules to force reimport
if "websocket._ssl_compat" in sys.modules:
del sys.modules["websocket._ssl_compat"]
# Mock the ssl module to not be available
import builtins
original_import = builtins.__import__
def mock_import(name, *args, **kwargs):
if name == "ssl":
raise ImportError("No module named 'ssl'")
return original_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=mock_import):
import websocket._ssl_compat as ssl_compat
# SSL should not be available
self.assertFalse(ssl_compat.HAVE_SSL)
self.assertIsNone(ssl_compat.ssl)
# Fallback exception classes should be available and functional
self.assertTrue(issubclass(ssl_compat.SSLError, Exception))
self.assertTrue(issubclass(ssl_compat.SSLEOFError, Exception))
self.assertTrue(issubclass(ssl_compat.SSLWantReadError, Exception))
self.assertTrue(issubclass(ssl_compat.SSLWantWriteError, Exception))
# Test that exceptions can be instantiated
ssl_error = ssl_compat.SSLError("test error")
self.assertIsInstance(ssl_error, Exception)
self.assertEqual(str(ssl_error), "test error")
ssl_eof_error = ssl_compat.SSLEOFError("test eof")
self.assertIsInstance(ssl_eof_error, Exception)
ssl_want_read = ssl_compat.SSLWantReadError("test read")
self.assertIsInstance(ssl_want_read, Exception)
ssl_want_write = ssl_compat.SSLWantWriteError("test write")
self.assertIsInstance(ssl_want_write, Exception)
def tearDown(self):
"""Clean up after tests"""
# Ensure ssl_compat is reimported fresh for next test
if "websocket._ssl_compat" in sys.modules:
del sys.modules["websocket._ssl_compat"]
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,638 @@
# -*- coding: utf-8 -*-
import unittest
import socket
import ssl
from unittest.mock import Mock, patch, MagicMock
from websocket._ssl_compat import (
SSLError,
SSLEOFError,
SSLWantReadError,
SSLWantWriteError,
HAVE_SSL,
)
from websocket._http import _ssl_socket, _wrap_sni_socket
from websocket._exceptions import WebSocketException
from websocket._socket import recv, send
"""
test_ssl_edge_cases.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class SSLEdgeCasesTest(unittest.TestCase):
def setUp(self):
if not HAVE_SSL:
self.skipTest("SSL not available")
def test_ssl_handshake_failure(self):
"""Test SSL handshake failure scenarios"""
mock_sock = Mock()
# Test SSL handshake timeout
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.side_effect = socket.timeout(
"SSL handshake timeout"
)
sslopt = {"cert_reqs": ssl.CERT_REQUIRED}
with self.assertRaises(socket.timeout):
_ssl_socket(mock_sock, sslopt, "example.com")
def test_ssl_certificate_verification_failures(self):
"""Test various SSL certificate verification failure scenarios"""
mock_sock = Mock()
# Test certificate verification failure
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.side_effect = ssl.SSLCertVerificationError(
"Certificate verification failed"
)
sslopt = {"cert_reqs": ssl.CERT_REQUIRED, "check_hostname": True}
with self.assertRaises(ssl.SSLCertVerificationError):
_ssl_socket(mock_sock, sslopt, "badssl.example")
def test_ssl_context_configuration_edge_cases(self):
"""Test SSL context configuration with various edge cases"""
mock_sock = Mock()
# Test with pre-created SSL context
with patch("ssl.SSLContext") as mock_ssl_context:
existing_context = Mock()
existing_context.wrap_socket.return_value = Mock()
mock_ssl_context.return_value = existing_context
sslopt = {"context": existing_context}
# Call _ssl_socket which should use the existing context
_ssl_socket(mock_sock, sslopt, "example.com")
# Should use the provided context, not create a new one
existing_context.wrap_socket.assert_called_once()
def test_ssl_ca_bundle_environment_edge_cases(self):
"""Test CA bundle environment variable edge cases"""
mock_sock = Mock()
# Test with non-existent CA bundle file
with patch.dict(
"os.environ", {"WEBSOCKET_CLIENT_CA_BUNDLE": "/nonexistent/ca-bundle.crt"}
):
with patch("os.path.isfile", return_value=False):
with patch("os.path.isdir", return_value=False):
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
sslopt = {}
_ssl_socket(mock_sock, sslopt, "example.com")
# Should not try to load non-existent CA bundle
mock_context.load_verify_locations.assert_not_called()
# Test with CA bundle directory
with patch.dict("os.environ", {"WEBSOCKET_CLIENT_CA_BUNDLE": "/etc/ssl/certs"}):
with patch("os.path.isfile", return_value=False):
with patch("os.path.isdir", return_value=True):
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
sslopt = {}
_ssl_socket(mock_sock, sslopt, "example.com")
# Should load CA directory
mock_context.load_verify_locations.assert_called_with(
cafile=None, capath="/etc/ssl/certs"
)
def test_ssl_cipher_configuration_edge_cases(self):
"""Test SSL cipher configuration edge cases"""
mock_sock = Mock()
# Test with invalid cipher suite
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.set_ciphers.side_effect = ssl.SSLError(
"No cipher can be selected"
)
mock_context.wrap_socket.return_value = Mock()
sslopt = {"ciphers": "INVALID_CIPHER"}
with self.assertRaises(WebSocketException):
_ssl_socket(mock_sock, sslopt, "example.com")
def test_ssl_ecdh_curve_edge_cases(self):
"""Test ECDH curve configuration edge cases"""
mock_sock = Mock()
# Test with invalid ECDH curve
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.set_ecdh_curve.side_effect = ValueError("unknown curve name")
mock_context.wrap_socket.return_value = Mock()
sslopt = {"ecdh_curve": "invalid_curve"}
with self.assertRaises(WebSocketException):
_ssl_socket(mock_sock, sslopt, "example.com")
def test_ssl_client_certificate_edge_cases(self):
"""Test client certificate configuration edge cases"""
mock_sock = Mock()
# Test with non-existent client certificate
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.load_cert_chain.side_effect = FileNotFoundError("No such file")
mock_context.wrap_socket.return_value = Mock()
sslopt = {"certfile": "/nonexistent/client.crt"}
with self.assertRaises(WebSocketException):
_ssl_socket(mock_sock, sslopt, "example.com")
def test_ssl_want_read_write_retry_edge_cases(self):
"""Test SSL want read/write retry edge cases"""
mock_sock = Mock()
# Test SSLWantReadError with multiple retries before success
read_attempts = [0] # Use list for mutable reference
def mock_recv(bufsize):
read_attempts[0] += 1
if read_attempts[0] == 1:
raise SSLWantReadError("The operation did not complete")
elif read_attempts[0] == 2:
return b"data after retries"
else:
return b""
mock_sock.recv.side_effect = mock_recv
mock_sock.gettimeout.return_value = 30.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Always ready
result = recv(mock_sock, 100)
self.assertEqual(result, b"data after retries")
self.assertEqual(read_attempts[0], 2)
# Should have used selector for retry
mock_selector.register.assert_called()
mock_selector.select.assert_called()
def test_ssl_want_write_retry_edge_cases(self):
"""Test SSL want write retry edge cases"""
mock_sock = Mock()
# Test SSLWantWriteError with multiple retries before success
write_attempts = [0] # Use list for mutable reference
def mock_send(data):
write_attempts[0] += 1
if write_attempts[0] == 1:
raise SSLWantWriteError("The operation did not complete")
elif write_attempts[0] == 2:
return len(data)
else:
return 0
mock_sock.send.side_effect = mock_send
mock_sock.gettimeout.return_value = 30.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True] # Always ready
result = send(mock_sock, b"test data")
self.assertEqual(result, 9) # len("test data")
self.assertEqual(write_attempts[0], 2)
def test_ssl_eof_error_edge_cases(self):
"""Test SSL EOF error edge cases"""
mock_sock = Mock()
# Test SSLEOFError during send
mock_sock.send.side_effect = SSLEOFError("SSL connection has been closed")
mock_sock.gettimeout.return_value = 30.0
from websocket._exceptions import WebSocketConnectionClosedException
with self.assertRaises(WebSocketConnectionClosedException):
send(mock_sock, b"test data")
def test_ssl_pending_data_edge_cases(self):
"""Test SSL pending data scenarios"""
from websocket._dispatcher import SSLDispatcher
from websocket._app import WebSocketApp
# Mock SSL socket with pending data
mock_ssl_sock = Mock()
mock_ssl_sock.pending.return_value = 1024 # Simulates pending SSL data
# Mock WebSocketApp
mock_app = Mock(spec=WebSocketApp)
mock_app.sock = Mock()
mock_app.sock.sock = mock_ssl_sock
dispatcher = SSLDispatcher(mock_app, 5.0)
# When there's pending data, should return immediately without selector
result = dispatcher.select(mock_ssl_sock, Mock())
# Should return the socket list when there's pending data
self.assertEqual(result, [mock_ssl_sock])
mock_ssl_sock.pending.assert_called_once()
def test_ssl_renegotiation_edge_cases(self):
"""Test SSL renegotiation scenarios"""
mock_sock = Mock()
# Simulate SSL renegotiation during read
call_count = 0
def mock_recv(bufsize):
nonlocal call_count
call_count += 1
if call_count == 1:
raise SSLWantReadError("SSL renegotiation required")
return b"data after renegotiation"
mock_sock.recv.side_effect = mock_recv
mock_sock.gettimeout.return_value = 30.0
with patch("selectors.DefaultSelector") as mock_selector_class:
mock_selector = Mock()
mock_selector_class.return_value = mock_selector
mock_selector.select.return_value = [True]
result = recv(mock_sock, 100)
self.assertEqual(result, b"data after renegotiation")
self.assertEqual(call_count, 2)
def test_ssl_server_hostname_override(self):
"""Test SSL server hostname override scenarios"""
mock_sock = Mock()
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
# Test server_hostname override
sslopt = {"server_hostname": "override.example.com"}
_ssl_socket(mock_sock, sslopt, "original.example.com")
# Should use override hostname in wrap_socket call
mock_context.wrap_socket.assert_called_with(
mock_sock,
do_handshake_on_connect=True,
suppress_ragged_eofs=True,
server_hostname="override.example.com",
)
def test_ssl_protocol_version_edge_cases(self):
"""Test SSL protocol version edge cases"""
mock_sock = Mock()
# Test with deprecated SSL version
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
# Test that deprecated ssl_version is still handled
if hasattr(ssl, "PROTOCOL_TLS"):
sslopt = {"ssl_version": ssl.PROTOCOL_TLS}
_ssl_socket(mock_sock, sslopt, "example.com")
mock_ssl_context.assert_called_with(ssl.PROTOCOL_TLS)
def test_ssl_keylog_file_edge_cases(self):
"""Test SSL keylog file configuration edge cases"""
mock_sock = Mock()
# Test with SSLKEYLOGFILE environment variable
with patch.dict("os.environ", {"SSLKEYLOGFILE": "/tmp/ssl_keys.log"}):
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
sslopt = {}
_ssl_socket(mock_sock, sslopt, "example.com")
# Should set keylog_filename
self.assertEqual(mock_context.keylog_filename, "/tmp/ssl_keys.log")
def test_ssl_context_verification_modes(self):
"""Test different SSL verification mode combinations"""
mock_sock = Mock()
test_cases = [
# (cert_reqs, check_hostname, expected_verify_mode, expected_check_hostname)
(ssl.CERT_NONE, False, ssl.CERT_NONE, False),
(ssl.CERT_REQUIRED, False, ssl.CERT_REQUIRED, False),
(ssl.CERT_REQUIRED, True, ssl.CERT_REQUIRED, True),
]
for cert_reqs, check_hostname, expected_verify, expected_check in test_cases:
with self.subTest(cert_reqs=cert_reqs, check_hostname=check_hostname):
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
sslopt = {"cert_reqs": cert_reqs, "check_hostname": check_hostname}
_ssl_socket(mock_sock, sslopt, "example.com")
self.assertEqual(mock_context.verify_mode, expected_verify)
self.assertEqual(mock_context.check_hostname, expected_check)
def test_ssl_socket_shutdown_edge_cases(self):
"""Test SSL socket shutdown edge cases"""
from websocket._core import WebSocket
mock_ssl_sock = Mock()
mock_ssl_sock.shutdown.side_effect = SSLError("SSL shutdown failed")
ws = WebSocket()
ws.sock = mock_ssl_sock
ws.connected = True
# Should handle SSL shutdown errors gracefully
try:
ws.close()
except SSLError:
self.fail("SSL shutdown error should be handled gracefully")
def test_ssl_socket_close_during_operation(self):
"""Test SSL socket being closed during ongoing operations"""
mock_sock = Mock()
# Simulate SSL socket being closed during recv
mock_sock.recv.side_effect = SSLError(
"SSL connection has been closed unexpectedly"
)
mock_sock.gettimeout.return_value = 30.0
from websocket._exceptions import WebSocketConnectionClosedException
# Should handle unexpected SSL closure
with self.assertRaises((SSLError, WebSocketConnectionClosedException)):
recv(mock_sock, 100)
def test_ssl_compression_edge_cases(self):
"""Test SSL compression configuration edge cases"""
mock_sock = Mock()
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
# Test SSL compression options (if available)
sslopt = {"compression": False} # Some SSL contexts support this
try:
_ssl_socket(mock_sock, sslopt, "example.com")
# Should not fail even if compression option is not supported
except AttributeError:
# Expected if SSL context doesn't support compression option
pass
def test_ssl_session_reuse_edge_cases(self):
"""Test SSL session reuse scenarios"""
mock_sock = Mock()
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_ssl_sock = Mock()
mock_context.wrap_socket.return_value = mock_ssl_sock
# Test session reuse
mock_ssl_sock.session = "mock_session"
mock_ssl_sock.session_reused = True
result = _ssl_socket(mock_sock, {}, "example.com")
# Should handle session reuse without issues
self.assertIsNotNone(result)
def test_ssl_alpn_protocol_edge_cases(self):
"""Test SSL ALPN (Application Layer Protocol Negotiation) edge cases"""
mock_sock = Mock()
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
# Test ALPN configuration
sslopt = {"alpn_protocols": ["http/1.1", "h2"]}
# ALPN protocols are not currently supported in the SSL wrapper
# but the test should not fail
result = _ssl_socket(mock_sock, sslopt, "example.com")
self.assertIsNotNone(result)
# ALPN would need to be implemented in _wrap_sni_socket function
def test_ssl_sni_edge_cases(self):
"""Test SSL SNI (Server Name Indication) edge cases"""
mock_sock = Mock()
# Test with IPv6 address (should not use SNI)
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
# IPv6 addresses should not be used for SNI
ipv6_hostname = "2001:db8::1"
_ssl_socket(mock_sock, {}, ipv6_hostname)
# Should use IPv6 address as server_hostname
mock_context.wrap_socket.assert_called_with(
mock_sock,
do_handshake_on_connect=True,
suppress_ragged_eofs=True,
server_hostname=ipv6_hostname,
)
def test_ssl_buffer_size_edge_cases(self):
"""Test SSL buffer size related edge cases"""
mock_sock = Mock()
def mock_recv(bufsize):
# SSL should never try to read more than 16KB at once
if bufsize > 16384:
raise SSLError("[SSL: BAD_LENGTH] buffer too large")
return b"A" * min(bufsize, 1024) # Return smaller chunks
mock_sock.recv.side_effect = mock_recv
mock_sock.gettimeout.return_value = 30.0
from websocket._abnf import frame_buffer
# Frame buffer should handle large requests by chunking
fb = frame_buffer(lambda size: recv(mock_sock, size), skip_utf8_validation=True)
# This should work even with large size due to chunking
result = fb.recv_strict(16384) # Exactly 16KB
self.assertGreater(len(result), 0)
def test_ssl_protocol_downgrade_protection(self):
"""Test SSL protocol downgrade protection"""
mock_sock = Mock()
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.side_effect = ssl.SSLError(
"SSLV3_ALERT_HANDSHAKE_FAILURE"
)
sslopt = {"ssl_version": ssl.PROTOCOL_TLS_CLIENT}
# Should propagate SSL protocol errors
with self.assertRaises(ssl.SSLError):
_ssl_socket(mock_sock, sslopt, "example.com")
def test_ssl_certificate_chain_validation(self):
"""Test SSL certificate chain validation edge cases"""
mock_sock = Mock()
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
# Test certificate chain validation failure
mock_context.wrap_socket.side_effect = ssl.SSLCertVerificationError(
"certificate verify failed: certificate has expired"
)
sslopt = {"cert_reqs": ssl.CERT_REQUIRED, "check_hostname": True}
with self.assertRaises(ssl.SSLCertVerificationError):
_ssl_socket(mock_sock, sslopt, "expired.badssl.com")
def test_ssl_weak_cipher_rejection(self):
"""Test SSL weak cipher rejection scenarios"""
mock_sock = Mock()
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.side_effect = ssl.SSLError("no shared cipher")
sslopt = {"ciphers": "RC4-MD5"} # Intentionally weak cipher
# Should fail with weak ciphers (SSL error is not wrapped by our code)
with self.assertRaises(ssl.SSLError):
_ssl_socket(mock_sock, sslopt, "example.com")
def test_ssl_hostname_verification_edge_cases(self):
"""Test SSL hostname verification edge cases"""
mock_sock = Mock()
# Test with wildcard certificate scenarios
test_cases = [
("*.example.com", "subdomain.example.com"), # Valid wildcard
("*.example.com", "sub.subdomain.example.com"), # Invalid wildcard depth
("example.com", "www.example.com"), # Hostname mismatch
]
for cert_hostname, connect_hostname in test_cases:
with self.subTest(cert=cert_hostname, hostname=connect_hostname):
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
if (
cert_hostname != connect_hostname
and "sub.subdomain" in connect_hostname
):
# Simulate hostname verification failure for invalid wildcard
mock_context.wrap_socket.side_effect = ssl.SSLCertVerificationError(
f"hostname '{connect_hostname}' doesn't match '{cert_hostname}'"
)
sslopt = {
"cert_reqs": ssl.CERT_REQUIRED,
"check_hostname": True,
}
with self.assertRaises(ssl.SSLCertVerificationError):
_ssl_socket(mock_sock, sslopt, connect_hostname)
else:
mock_context.wrap_socket.return_value = Mock()
sslopt = {
"cert_reqs": ssl.CERT_REQUIRED,
"check_hostname": True,
}
# Should succeed for valid cases
result = _ssl_socket(mock_sock, sslopt, connect_hostname)
self.assertIsNotNone(result)
def test_ssl_memory_bio_edge_cases(self):
"""Test SSL memory BIO edge cases"""
mock_sock = Mock()
# Test SSL memory BIO scenarios (if available)
try:
import ssl
if hasattr(ssl, "MemoryBIO"):
with patch("ssl.SSLContext") as mock_ssl_context:
mock_context = Mock()
mock_ssl_context.return_value = mock_context
mock_context.wrap_socket.return_value = Mock()
# Memory BIO should work if available
_ssl_socket(mock_sock, {}, "example.com")
# Standard socket wrapping should still work
mock_context.wrap_socket.assert_called_once()
except (ImportError, AttributeError):
self.skipTest("SSL MemoryBIO not available")
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,471 @@
# -*- coding: utf-8 -*-
#
import os
import unittest
from websocket._url import (
_is_address_in_network,
_is_no_proxy_host,
get_proxy_info,
parse_url,
)
from websocket._exceptions import WebSocketProxyException
"""
test_url.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class UrlTest(unittest.TestCase):
def test_address_in_network(self):
self.assertTrue(_is_address_in_network("127.0.0.1", "127.0.0.0/8"))
self.assertTrue(_is_address_in_network("127.1.0.1", "127.0.0.0/8"))
self.assertFalse(_is_address_in_network("127.1.0.1", "127.0.0.0/24"))
self.assertTrue(_is_address_in_network("2001:db8::1", "2001:db8::/64"))
self.assertFalse(_is_address_in_network("2001:db8:1::1", "2001:db8::/64"))
def test_parse_url(self):
p = parse_url("ws://www.example.com/r")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 80)
self.assertEqual(p[2], "/r")
self.assertEqual(p[3], False)
p = parse_url("ws://www.example.com/r/")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 80)
self.assertEqual(p[2], "/r/")
self.assertEqual(p[3], False)
p = parse_url("ws://www.example.com/")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 80)
self.assertEqual(p[2], "/")
self.assertEqual(p[3], False)
p = parse_url("ws://www.example.com")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 80)
self.assertEqual(p[2], "/")
self.assertEqual(p[3], False)
p = parse_url("ws://www.example.com:8080/r")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 8080)
self.assertEqual(p[2], "/r")
self.assertEqual(p[3], False)
p = parse_url("ws://www.example.com:8080/")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 8080)
self.assertEqual(p[2], "/")
self.assertEqual(p[3], False)
p = parse_url("ws://www.example.com:8080")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 8080)
self.assertEqual(p[2], "/")
self.assertEqual(p[3], False)
p = parse_url("wss://www.example.com:8080/r")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 8080)
self.assertEqual(p[2], "/r")
self.assertEqual(p[3], True)
p = parse_url("wss://www.example.com:8080/r?key=value")
self.assertEqual(p[0], "www.example.com")
self.assertEqual(p[1], 8080)
self.assertEqual(p[2], "/r?key=value")
self.assertEqual(p[3], True)
self.assertRaises(ValueError, parse_url, "http://www.example.com/r")
p = parse_url("ws://[2a03:4000:123:83::3]/r")
self.assertEqual(p[0], "2a03:4000:123:83::3")
self.assertEqual(p[1], 80)
self.assertEqual(p[2], "/r")
self.assertEqual(p[3], False)
p = parse_url("ws://[2a03:4000:123:83::3]:8080/r")
self.assertEqual(p[0], "2a03:4000:123:83::3")
self.assertEqual(p[1], 8080)
self.assertEqual(p[2], "/r")
self.assertEqual(p[3], False)
p = parse_url("wss://[2a03:4000:123:83::3]/r")
self.assertEqual(p[0], "2a03:4000:123:83::3")
self.assertEqual(p[1], 443)
self.assertEqual(p[2], "/r")
self.assertEqual(p[3], True)
p = parse_url("wss://[2a03:4000:123:83::3]:8080/r")
self.assertEqual(p[0], "2a03:4000:123:83::3")
self.assertEqual(p[1], 8080)
self.assertEqual(p[2], "/r")
self.assertEqual(p[3], True)
class IsNoProxyHostTest(unittest.TestCase):
def setUp(self):
self.no_proxy = os.environ.get("no_proxy", None)
if "no_proxy" in os.environ:
del os.environ["no_proxy"]
def tearDown(self):
if self.no_proxy:
os.environ["no_proxy"] = self.no_proxy
elif "no_proxy" in os.environ:
del os.environ["no_proxy"]
def test_match_all(self):
self.assertTrue(_is_no_proxy_host("any.websocket.org", ["*"]))
self.assertTrue(_is_no_proxy_host("192.168.0.1", ["*"]))
self.assertFalse(_is_no_proxy_host("192.168.0.1", ["192.168.1.1"]))
self.assertFalse(
_is_no_proxy_host("any.websocket.org", ["other.websocket.org"])
)
self.assertTrue(
_is_no_proxy_host("any.websocket.org", ["other.websocket.org", "*"])
)
os.environ["no_proxy"] = "*"
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
self.assertTrue(_is_no_proxy_host("192.168.0.1", None))
os.environ["no_proxy"] = "other.websocket.org, *"
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
def test_ip_address(self):
self.assertTrue(_is_no_proxy_host("127.0.0.1", ["127.0.0.1"]))
self.assertFalse(_is_no_proxy_host("127.0.0.2", ["127.0.0.1"]))
self.assertTrue(
_is_no_proxy_host("127.0.0.1", ["other.websocket.org", "127.0.0.1"])
)
self.assertFalse(
_is_no_proxy_host("127.0.0.2", ["other.websocket.org", "127.0.0.1"])
)
os.environ["no_proxy"] = "127.0.0.1"
self.assertTrue(_is_no_proxy_host("127.0.0.1", None))
self.assertFalse(_is_no_proxy_host("127.0.0.2", None))
os.environ["no_proxy"] = "other.websocket.org, 127.0.0.1"
self.assertTrue(_is_no_proxy_host("127.0.0.1", None))
self.assertFalse(_is_no_proxy_host("127.0.0.2", None))
def test_ip_address_in_range(self):
self.assertTrue(_is_no_proxy_host("127.0.0.1", ["127.0.0.0/8"]))
self.assertTrue(_is_no_proxy_host("127.0.0.2", ["127.0.0.0/8"]))
self.assertFalse(_is_no_proxy_host("127.1.0.1", ["127.0.0.0/24"]))
self.assertTrue(_is_no_proxy_host("2001:db8::1", ["2001:db8::/64"]))
self.assertFalse(_is_no_proxy_host("2001:db8:1::1", ["2001:db8::/64"]))
os.environ["no_proxy"] = "127.0.0.0/8,2001:db8::/64"
self.assertTrue(_is_no_proxy_host("127.0.0.1", None))
self.assertTrue(_is_no_proxy_host("127.0.0.2", None))
self.assertTrue(_is_no_proxy_host("2001:db8::1", None))
self.assertFalse(_is_no_proxy_host("2001:db8:1::1", None))
os.environ["no_proxy"] = "127.0.0.0/24,2001:db8::/64"
self.assertFalse(_is_no_proxy_host("127.1.0.1", None))
self.assertFalse(_is_no_proxy_host("2001:db8:1::1", None))
def test_hostname_match(self):
self.assertTrue(_is_no_proxy_host("my.websocket.org", ["my.websocket.org"]))
self.assertTrue(
_is_no_proxy_host(
"my.websocket.org", ["other.websocket.org", "my.websocket.org"]
)
)
self.assertFalse(_is_no_proxy_host("my.websocket.org", ["other.websocket.org"]))
os.environ["no_proxy"] = "my.websocket.org"
self.assertTrue(_is_no_proxy_host("my.websocket.org", None))
self.assertFalse(_is_no_proxy_host("other.websocket.org", None))
os.environ["no_proxy"] = "other.websocket.org, my.websocket.org"
self.assertTrue(_is_no_proxy_host("my.websocket.org", None))
def test_hostname_match_domain(self):
self.assertTrue(_is_no_proxy_host("any.websocket.org", [".websocket.org"]))
self.assertTrue(_is_no_proxy_host("my.other.websocket.org", [".websocket.org"]))
self.assertTrue(
_is_no_proxy_host(
"any.websocket.org", ["my.websocket.org", ".websocket.org"]
)
)
self.assertFalse(_is_no_proxy_host("any.websocket.com", [".websocket.org"]))
os.environ["no_proxy"] = ".websocket.org"
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
self.assertTrue(_is_no_proxy_host("my.other.websocket.org", None))
self.assertFalse(_is_no_proxy_host("any.websocket.com", None))
os.environ["no_proxy"] = "my.websocket.org, .websocket.org"
self.assertTrue(_is_no_proxy_host("any.websocket.org", None))
class ProxyInfoTest(unittest.TestCase):
def setUp(self):
self.http_proxy = os.environ.get("http_proxy", None)
self.https_proxy = os.environ.get("https_proxy", None)
self.no_proxy = os.environ.get("no_proxy", None)
if "http_proxy" in os.environ:
del os.environ["http_proxy"]
if "https_proxy" in os.environ:
del os.environ["https_proxy"]
if "no_proxy" in os.environ:
del os.environ["no_proxy"]
def tearDown(self):
if self.http_proxy:
os.environ["http_proxy"] = self.http_proxy
elif "http_proxy" in os.environ:
del os.environ["http_proxy"]
if self.https_proxy:
os.environ["https_proxy"] = self.https_proxy
elif "https_proxy" in os.environ:
del os.environ["https_proxy"]
if self.no_proxy:
os.environ["no_proxy"] = self.no_proxy
elif "no_proxy" in os.environ:
del os.environ["no_proxy"]
def test_proxy_from_args(self):
self.assertRaises(
WebSocketProxyException,
get_proxy_info,
"echo.websocket.events",
False,
proxy_host="localhost",
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events", False, proxy_host="localhost", proxy_port=3128
),
("localhost", 3128, None),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events", True, proxy_host="localhost", proxy_port=3128
),
("localhost", 3128, None),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events",
False,
proxy_host="localhost",
proxy_port=9001,
proxy_auth=("a", "b"),
),
("localhost", 9001, ("a", "b")),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events",
False,
proxy_host="localhost",
proxy_port=3128,
proxy_auth=("a", "b"),
),
("localhost", 3128, ("a", "b")),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events",
True,
proxy_host="localhost",
proxy_port=8765,
proxy_auth=("a", "b"),
),
("localhost", 8765, ("a", "b")),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events",
True,
proxy_host="localhost",
proxy_port=3128,
proxy_auth=("a", "b"),
),
("localhost", 3128, ("a", "b")),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events",
True,
proxy_host="localhost",
proxy_port=3128,
no_proxy=["example.com"],
proxy_auth=("a", "b"),
),
("localhost", 3128, ("a", "b")),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events",
True,
proxy_host="localhost",
proxy_port=3128,
no_proxy=["echo.websocket.events"],
proxy_auth=("a", "b"),
),
(None, 0, None),
)
self.assertEqual(
get_proxy_info(
"echo.websocket.events",
True,
proxy_host="localhost",
proxy_port=3128,
no_proxy=[".websocket.events"],
),
(None, 0, None),
)
def test_proxy_from_env(self):
os.environ["http_proxy"] = "http://localhost/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
)
os.environ["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = "http://localhost2/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = "http://localhost2:3128/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
)
os.environ["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = "http://localhost2/"
self.assertEqual(
get_proxy_info("echo.websocket.events", True), ("localhost2", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = "http://localhost2:3128/"
self.assertEqual(
get_proxy_info("echo.websocket.events", True), ("localhost2", 3128, None)
)
os.environ["http_proxy"] = ""
os.environ["https_proxy"] = "http://localhost2/"
self.assertEqual(
get_proxy_info("echo.websocket.events", True), ("localhost2", None, None)
)
self.assertEqual(
get_proxy_info("echo.websocket.events", False), (None, 0, None)
)
os.environ["http_proxy"] = ""
os.environ["https_proxy"] = "http://localhost2:3128/"
self.assertEqual(
get_proxy_info("echo.websocket.events", True), ("localhost2", 3128, None)
)
self.assertEqual(
get_proxy_info("echo.websocket.events", False), (None, 0, None)
)
os.environ["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = ""
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", None, None)
)
os.environ["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = ""
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
self.assertEqual(
get_proxy_info("echo.websocket.events", False), ("localhost", 3128, None)
)
os.environ["http_proxy"] = "http://a:b@localhost/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", None, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", 3128, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost/"
os.environ["https_proxy"] = "http://a:b@localhost2/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", None, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
self.assertEqual(
get_proxy_info("echo.websocket.events", False),
("localhost", 3128, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost/"
os.environ["https_proxy"] = "http://a:b@localhost2/"
self.assertEqual(
get_proxy_info("echo.websocket.events", True),
("localhost2", None, ("a", "b")),
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
self.assertEqual(
get_proxy_info("echo.websocket.events", True),
("localhost2", 3128, ("a", "b")),
)
os.environ["http_proxy"] = (
"http://john%40example.com:P%40SSWORD@localhost:3128/"
)
os.environ["https_proxy"] = (
"http://john%40example.com:P%40SSWORD@localhost2:3128/"
)
self.assertEqual(
get_proxy_info("echo.websocket.events", True),
("localhost2", 3128, ("john@example.com", "P@SSWORD")),
)
os.environ["http_proxy"] = "http://a:b@localhost/"
os.environ["https_proxy"] = "http://a:b@localhost2/"
os.environ["no_proxy"] = "example1.com,example2.com"
self.assertEqual(
get_proxy_info("example.1.com", True), ("localhost2", None, ("a", "b"))
)
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
os.environ["no_proxy"] = "example1.com,example2.com, echo.websocket.events"
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
os.environ["no_proxy"] = "example1.com,example2.com, .websocket.events"
self.assertEqual(get_proxy_info("echo.websocket.events", True), (None, 0, None))
os.environ["http_proxy"] = "http://a:b@localhost:3128/"
os.environ["https_proxy"] = "http://a:b@localhost2:3128/"
os.environ["no_proxy"] = "127.0.0.0/8, 192.168.0.0/16"
self.assertEqual(get_proxy_info("127.0.0.1", False), (None, 0, None))
self.assertEqual(get_proxy_info("192.168.1.1", False), (None, 0, None))
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,138 @@
# -*- coding: utf-8 -*-
import sys
import unittest
from unittest.mock import patch
"""
test_utils.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
class UtilsTest(unittest.TestCase):
def test_nolock(self):
"""Test NoLock context manager"""
from websocket._utils import NoLock
lock = NoLock()
# Test that it can be used as context manager
with lock:
pass # Should not raise any exception
# Test enter/exit methods directly
self.assertIsNone(lock.__enter__())
self.assertIsNone(lock.__exit__(None, None, None))
def test_utf8_validation_with_wsaccel(self):
"""Test UTF-8 validation when wsaccel is available"""
# Import normally (wsaccel should be available in test environment)
from websocket._utils import validate_utf8
# Test valid UTF-8 strings (convert to bytes for wsaccel)
self.assertTrue(validate_utf8("Hello, World!".encode("utf-8")))
self.assertTrue(validate_utf8("🌟 Unicode test".encode("utf-8")))
self.assertTrue(validate_utf8(b"Hello, bytes"))
self.assertTrue(validate_utf8("Héllo with accénts".encode("utf-8")))
# Test invalid UTF-8 sequences
self.assertFalse(validate_utf8(b"\xff\xfe")) # Invalid UTF-8
self.assertFalse(validate_utf8(b"\x80\x80")) # Invalid continuation
def test_utf8_validation_fallback(self):
"""Test UTF-8 validation fallback when wsaccel is not available"""
# Remove _utils from modules to force reimport
if "websocket._utils" in sys.modules:
del sys.modules["websocket._utils"]
# Mock wsaccel import to raise ImportError
import builtins
original_import = builtins.__import__
def mock_import(name, *args, **kwargs):
if "wsaccel" in name:
raise ImportError(f"No module named '{name}'")
return original_import(name, *args, **kwargs)
with patch("builtins.__import__", side_effect=mock_import):
import websocket._utils as utils
# Test valid UTF-8 strings with fallback implementation (convert strings to bytes)
self.assertTrue(utils.validate_utf8("Hello, World!".encode("utf-8")))
self.assertTrue(utils.validate_utf8(b"Hello, bytes"))
self.assertTrue(utils.validate_utf8("ASCII text".encode("utf-8")))
# Test Unicode strings (convert to bytes)
self.assertTrue(utils.validate_utf8("🌟 Unicode test".encode("utf-8")))
self.assertTrue(utils.validate_utf8("Héllo with accénts".encode("utf-8")))
# Test empty string/bytes
self.assertTrue(utils.validate_utf8("".encode("utf-8")))
self.assertTrue(utils.validate_utf8(b""))
# Test invalid UTF-8 sequences (should return False)
self.assertFalse(utils.validate_utf8(b"\xff\xfe"))
self.assertFalse(utils.validate_utf8(b"\x80\x80"))
# Note: The fallback implementation may have different validation behavior
# than wsaccel, so we focus on clearly invalid sequences
def test_extract_err_message(self):
"""Test extract_err_message function"""
from websocket._utils import extract_err_message
# Test with exception that has args
exc_with_args = Exception("Test error message")
self.assertEqual(extract_err_message(exc_with_args), "Test error message")
# Test with exception that has multiple args
exc_multi_args = Exception("First arg", "Second arg")
self.assertEqual(extract_err_message(exc_multi_args), "First arg")
# Test with exception that has no args
exc_no_args = Exception()
self.assertIsNone(extract_err_message(exc_no_args))
def test_extract_error_code(self):
"""Test extract_error_code function"""
from websocket._utils import extract_error_code
# Test with exception that has integer as first arg
exc_with_code = Exception(404, "Not found")
self.assertEqual(extract_error_code(exc_with_code), 404)
# Test with exception that has string as first arg
exc_with_string = Exception("Error message", "Second arg")
self.assertIsNone(extract_error_code(exc_with_string))
# Test with exception that has only one arg
exc_single_arg = Exception("Single arg")
self.assertIsNone(extract_error_code(exc_single_arg))
# Test with exception that has no args
exc_no_args = Exception()
self.assertIsNone(extract_error_code(exc_no_args))
def tearDown(self):
"""Clean up after tests"""
# Ensure _utils is reimported fresh for next test
if "websocket._utils" in sys.modules:
del sys.modules["websocket._utils"]
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,501 @@
# -*- coding: utf-8 -*-
#
import os
import os.path
import socket
import unittest
from base64 import decodebytes as base64decode
import websocket as ws
from websocket._exceptions import (
WebSocketBadStatusException,
WebSocketAddressException,
WebSocketException,
)
from websocket._handshake import _create_sec_websocket_key
from websocket._handshake import _validate as _validate_header
from websocket._http import read_headers
from websocket._utils import validate_utf8
"""
test_websocket.py
websocket - WebSocket client library for Python
Copyright 2025 engn33r
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
try:
import ssl
except ImportError:
# dummy class of SSLError for ssl none-support environment.
class SSLError(Exception):
pass
# Skip test to access the internet unless TEST_WITH_INTERNET == 1
TEST_WITH_INTERNET = os.environ.get("TEST_WITH_INTERNET", "0") == "1"
# Skip tests relying on local websockets server unless LOCAL_WS_SERVER_PORT != -1
LOCAL_WS_SERVER_PORT = os.environ.get("LOCAL_WS_SERVER_PORT", "-1")
TEST_WITH_LOCAL_SERVER = LOCAL_WS_SERVER_PORT != "-1"
TRACEABLE = True
def create_mask_key(_):
return "abcd"
class SockMock:
def __init__(self):
self.data = []
self.sent = []
def add_packet(self, data):
self.data.append(data)
def gettimeout(self):
return None
def recv(self, bufsize):
if self.data:
e = self.data.pop(0)
if isinstance(e, Exception):
raise e
if len(e) > bufsize:
self.data.insert(0, e[bufsize:])
return e[:bufsize]
def send(self, data):
self.sent.append(data)
return len(data)
def close(self):
pass
class HeaderSockMock(SockMock):
def __init__(self, fname):
SockMock.__init__(self)
path = os.path.join(os.path.dirname(__file__), fname)
with open(path, "rb") as f:
self.add_packet(f.read())
class WebSocketTest(unittest.TestCase):
def setUp(self):
ws.enableTrace(TRACEABLE)
def tearDown(self):
pass
def test_default_timeout(self):
self.assertEqual(ws.getdefaulttimeout(), None)
ws.setdefaulttimeout(10)
self.assertEqual(ws.getdefaulttimeout(), 10)
ws.setdefaulttimeout(None)
def test_ws_key(self):
key = _create_sec_websocket_key()
self.assertTrue(key != 24)
self.assertTrue("¥n" not in key)
def test_nonce(self):
"""WebSocket key should be a random 16-byte nonce."""
key = _create_sec_websocket_key()
nonce = base64decode(key.encode("utf-8"))
self.assertEqual(16, len(nonce))
def test_ws_utils(self):
key = "c6b8hTg4EeGb2gQMztV1/g=="
required_header = {
"upgrade": "websocket",
"connection": "upgrade",
"sec-websocket-accept": "Kxep+hNu9n51529fGidYu7a3wO0=",
}
self.assertEqual(_validate_header(required_header, key, None), (True, None))
header = required_header.copy()
header["upgrade"] = "http"
self.assertEqual(_validate_header(header, key, None), (False, None))
del header["upgrade"]
self.assertEqual(_validate_header(header, key, None), (False, None))
header = required_header.copy()
header["connection"] = "something"
self.assertEqual(_validate_header(header, key, None), (False, None))
del header["connection"]
self.assertEqual(_validate_header(header, key, None), (False, None))
header = required_header.copy()
header["sec-websocket-accept"] = "something"
self.assertEqual(_validate_header(header, key, None), (False, None))
del header["sec-websocket-accept"]
self.assertEqual(_validate_header(header, key, None), (False, None))
header = required_header.copy()
header["sec-websocket-protocol"] = "sub1"
self.assertEqual(
_validate_header(header, key, ["sub1", "sub2"]), (True, "sub1")
)
# This case will print out a logging error using the error() function, but that is expected
self.assertEqual(_validate_header(header, key, ["sub2", "sub3"]), (False, None))
header = required_header.copy()
header["sec-websocket-protocol"] = "sUb1"
self.assertEqual(
_validate_header(header, key, ["Sub1", "suB2"]), (True, "sub1")
)
header = required_header.copy()
# This case will print out a logging error using the error() function, but that is expected
self.assertEqual(_validate_header(header, key, ["Sub1", "suB2"]), (False, None))
def test_read_header(self):
status, header, _ = read_headers(HeaderSockMock("data/header01.txt"))
self.assertEqual(status, 101)
self.assertEqual(header["connection"], "Upgrade")
status, header, _ = read_headers(HeaderSockMock("data/header03.txt"))
self.assertEqual(status, 101)
self.assertEqual(header["connection"], "Upgrade, Keep-Alive")
HeaderSockMock("data/header02.txt")
self.assertRaises(
ws.WebSocketException, read_headers, HeaderSockMock("data/header02.txt")
)
def test_send(self):
# TODO: add longer frame data
sock = ws.WebSocket()
sock.set_mask_key(create_mask_key)
s = sock.sock = HeaderSockMock("data/header01.txt")
sock.send("Hello")
self.assertEqual(s.sent[0], b"\x81\x85abcd)\x07\x0f\x08\x0e")
sock.send("こんにちは")
self.assertEqual(
s.sent[1],
b"\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc",
)
# sock.send("x" * 5000)
# self.assertEqual(s.sent[1], b'\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")
self.assertEqual(sock.send_binary(b"1111111111101"), 19)
def test_recv(self):
# TODO: add longer frame data
sock = ws.WebSocket()
s = sock.sock = SockMock()
something = (
b"\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc"
)
s.add_packet(something)
data = sock.recv()
self.assertEqual(data, "こんにちは")
s.add_packet(b"\x81\x85abcd)\x07\x0f\x08\x0e")
data = sock.recv()
self.assertEqual(data, "Hello")
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_iter(self):
count = 2
s = ws.create_connection("wss://api.bitfinex.com/ws/2")
s.send('{"event": "subscribe", "channel": "ticker"}')
for _ in s:
count -= 1
if count == 0:
break
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_next(self):
sock = ws.create_connection("wss://api.bitfinex.com/ws/2")
self.assertEqual(str, type(next(sock)))
def test_internal_recv_strict(self):
sock = ws.WebSocket()
s = sock.sock = SockMock()
s.add_packet(b"foo")
s.add_packet(socket.timeout())
s.add_packet(b"bar")
# s.add_packet(SSLError("The read operation timed out"))
s.add_packet(b"baz")
with self.assertRaises(ws.WebSocketTimeoutException):
sock.frame_buffer.recv_strict(9)
# with self.assertRaises(SSLError):
# data = sock._recv_strict(9)
data = sock.frame_buffer.recv_strict(9)
self.assertEqual(data, b"foobarbaz")
with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.frame_buffer.recv_strict(1)
def test_recv_timeout(self):
sock = ws.WebSocket()
s = sock.sock = SockMock()
s.add_packet(b"\x81")
s.add_packet(socket.timeout())
s.add_packet(b"\x8dabcd\x29\x07\x0f\x08\x0e")
s.add_packet(socket.timeout())
s.add_packet(b"\x4e\x43\x33\x0e\x10\x0f\x00\x40")
with self.assertRaises(ws.WebSocketTimeoutException):
sock.recv()
with self.assertRaises(ws.WebSocketTimeoutException):
sock.recv()
data = sock.recv()
self.assertEqual(data, "Hello, World!")
with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.recv()
def test_recv_with_simple_fragmentation(self):
sock = ws.WebSocket()
s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Brevity is "
s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
# OPCODE=CONT, FIN=1, MSG="the soul of wit"
s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
data = sock.recv()
self.assertEqual(data, "Brevity is the soul of wit")
with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.recv()
def test_recv_with_fire_event_of_fragmentation(self):
sock = ws.WebSocket(fire_cont_frame=True)
s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Brevity is "
s.add_packet(b"\x01\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
# OPCODE=CONT, FIN=0, MSG="Brevity is "
s.add_packet(b"\x00\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
# OPCODE=CONT, FIN=1, MSG="the soul of wit"
s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
_, data = sock.recv_data()
self.assertEqual(data, b"Brevity is ")
_, data = sock.recv_data()
self.assertEqual(data, b"Brevity is ")
_, data = sock.recv_data()
self.assertEqual(data, b"the soul of wit")
# OPCODE=CONT, FIN=0, MSG="Brevity is "
s.add_packet(b"\x80\x8babcd#\x10\x06\x12\x08\x16\x1aD\x08\x11C")
with self.assertRaises(ws.WebSocketException):
sock.recv_data()
with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.recv()
def test_close(self):
sock = ws.WebSocket()
sock.connected = True
sock.close()
sock = ws.WebSocket()
s = sock.sock = SockMock()
sock.connected = True
s.add_packet(b"\x88\x80\x17\x98p\x84")
sock.recv()
self.assertEqual(sock.connected, False)
def test_recv_cont_fragmentation(self):
sock = ws.WebSocket()
s = sock.sock = SockMock()
# OPCODE=CONT, FIN=1, MSG="the soul of wit"
s.add_packet(b"\x80\x8fabcd\x15\n\x06D\x12\r\x16\x08A\r\x05D\x16\x0b\x17")
self.assertRaises(ws.WebSocketException, sock.recv)
def test_recv_with_prolonged_fragmentation(self):
sock = ws.WebSocket()
s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Once more unto the breach, "
s.add_packet(
b"\x01\x9babcd.\x0c\x00\x01A\x0f\x0c\x16\x04B\x16\n\x15\rC\x10\t\x07C\x06\x13\x07\x02\x07\tNC"
)
# OPCODE=CONT, FIN=0, MSG="dear friends, "
s.add_packet(b"\x00\x8eabcd\x05\x07\x02\x16A\x04\x11\r\x04\x0c\x07\x17MB")
# OPCODE=CONT, FIN=1, MSG="once more"
s.add_packet(b"\x80\x89abcd\x0e\x0c\x00\x01A\x0f\x0c\x16\x04")
data = sock.recv()
self.assertEqual(data, "Once more unto the breach, dear friends, once more")
with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.recv()
def test_recv_with_fragmentation_and_control_frame(self):
sock = ws.WebSocket()
sock.set_mask_key(create_mask_key)
s = sock.sock = SockMock()
# OPCODE=TEXT, FIN=0, MSG="Too much "
s.add_packet(b"\x01\x89abcd5\r\x0cD\x0c\x17\x00\x0cA")
# OPCODE=PING, FIN=1, MSG="Please PONG this"
s.add_packet(b"\x89\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17")
# OPCODE=CONT, FIN=1, MSG="of a good thing"
s.add_packet(b"\x80\x8fabcd\x0e\x04C\x05A\x05\x0c\x0b\x05B\x17\x0c\x08\x0c\x04")
data = sock.recv()
self.assertEqual(data, "Too much of a good thing")
with self.assertRaises(ws.WebSocketConnectionClosedException):
sock.recv()
self.assertEqual(
s.sent[0], b"\x8a\x90abcd1\x0e\x06\x05\x12\x07C4.,$D\x15\n\n\x17"
)
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_websocket(self):
s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
self.assertNotEqual(s, None)
s.send("Hello, World")
result = s.next()
s.fileno()
self.assertEqual(result, "Hello, World")
s.send("こにゃにゃちは、世界")
result = s.recv()
self.assertEqual(result, "こにゃにゃちは、世界")
self.assertRaises(ValueError, s.send_close, -1, "")
s.close()
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_ping_pong(self):
s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
self.assertNotEqual(s, None)
s.ping("Hello")
s.pong("Hi")
s.close()
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_support_redirect(self):
s = ws.WebSocket()
self.assertRaises(WebSocketBadStatusException, s.connect, "ws://google.com/")
# Need to find a URL that has a redirect code leading to a websocket
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_secure_websocket(self):
s = ws.create_connection("wss://api.bitfinex.com/ws/2")
self.assertNotEqual(s, None)
self.assertTrue(isinstance(s.sock, ssl.SSLSocket))
self.assertEqual(s.getstatus(), 101)
self.assertNotEqual(s.getheaders(), None)
s.settimeout(10)
self.assertEqual(s.gettimeout(), 10)
self.assertEqual(s.getsubprotocol(), None)
s.abort()
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_websocket_with_custom_header(self):
s = ws.create_connection(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}",
headers={"User-Agent": "PythonWebsocketClient"},
)
self.assertNotEqual(s, None)
self.assertEqual(s.getsubprotocol(), None)
s.send("Hello, World")
result = s.recv()
self.assertEqual(result, "Hello, World")
self.assertRaises(ValueError, s.close, -1, "")
s.close()
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_after_close(self):
s = ws.create_connection(f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}")
self.assertNotEqual(s, None)
s.close()
self.assertRaises(ws.WebSocketConnectionClosedException, s.send, "Hello")
self.assertRaises(ws.WebSocketConnectionClosedException, s.recv)
class SockOptTest(unittest.TestCase):
@unittest.skipUnless(
TEST_WITH_LOCAL_SERVER, "Tests using local websocket server are disabled"
)
def test_sockopt(self):
sockopt = ((socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),)
s = ws.create_connection(
f"ws://127.0.0.1:{LOCAL_WS_SERVER_PORT}", sockopt=sockopt
)
self.assertNotEqual(
s.sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY), 0
)
s.close()
class UtilsTest(unittest.TestCase):
def test_utf8_validator(self):
state = validate_utf8(b"\xf0\x90\x80\x80")
self.assertEqual(state, True)
state = validate_utf8(
b"\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80edited"
)
self.assertEqual(state, False)
state = validate_utf8(b"")
self.assertEqual(state, True)
class HandshakeTest(unittest.TestCase):
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_http_ssl(self):
websock1 = ws.WebSocket(
sslopt={"cert_chain": ssl.get_default_verify_paths().capath},
enable_multithread=False,
)
self.assertRaises(ValueError, websock1.connect, "wss://api.bitfinex.com/ws/2")
websock2 = ws.WebSocket(sslopt={"certfile": "myNonexistentCertFile"})
self.assertRaises(
WebSocketException, websock2.connect, "wss://api.bitfinex.com/ws/2"
)
@unittest.skipUnless(TEST_WITH_INTERNET, "Internet-requiring tests are disabled")
def test_manual_headers(self):
websock3 = ws.WebSocket(
sslopt={
"ca_certs": ssl.get_default_verify_paths().cafile,
"ca_cert_path": ssl.get_default_verify_paths().capath,
}
)
self.assertRaises(
WebSocketBadStatusException,
websock3.connect,
"wss://api.bitfinex.com/ws/2",
cookie="chocolate",
origin="testing_websockets.com",
host="echo.websocket.events/websocket-client-test",
subprotocols=["testproto"],
connection="Upgrade",
header={
"CustomHeader1": "123",
"Cookie": "TestValue",
"Sec-WebSocket-Key": "k9kFAUWNAMmf5OEMfTlOEA==",
"Sec-WebSocket-Protocol": "newprotocol",
},
)
def test_ipv6(self):
websock2 = ws.WebSocket()
self.assertRaises(ValueError, websock2.connect, "2001:4860:4860::8888")
def test_bad_urls(self):
websock3 = ws.WebSocket()
self.assertRaises(ValueError, websock3.connect, "ws//example.com")
self.assertRaises(WebSocketAddressException, websock3.connect, "ws://example")
self.assertRaises(ValueError, websock3.connect, "example.com")
if __name__ == "__main__":
unittest.main()