161 lines
6.3 KiB
Python
161 lines
6.3 KiB
Python
# -*- 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()
|