codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2import errno3import socket4import unittest5from unittest.mock import Mock, patch, MagicMock6import time7 8from websocket._socket import recv, recv_line, send, DEFAULT_SOCKET_OPTION9from websocket._ssl_compat import (10    SSLError,11    SSLEOFError,12    SSLWantWriteError,13    SSLWantReadError,14)15from websocket._exceptions import (16    WebSocketTimeoutException,17    WebSocketConnectionClosedException,18)19 20"""21test_socket.py22websocket - WebSocket client library for Python23 24Copyright 2025 engn33r25 26Licensed under the Apache License, Version 2.0 (the "License");27you may not use this file except in compliance with the License.28You may obtain a copy of the License at29 30    http://www.apache.org/licenses/LICENSE-2.031 32Unless required by applicable law or agreed to in writing, software33distributed under the License is distributed on an "AS IS" BASIS,34WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.35See the License for the specific language governing permissions and36limitations under the License.37"""38 39class SocketTest(unittest.TestCase):40    def test_default_socket_option(self):41        """Test DEFAULT_SOCKET_OPTION contains expected options"""42        self.assertIsInstance(DEFAULT_SOCKET_OPTION, list)43        self.assertGreater(len(DEFAULT_SOCKET_OPTION), 0)44 45        # Should contain TCP_NODELAY option46        tcp_nodelay_found = any(47            opt[1] == socket.TCP_NODELAY for opt in DEFAULT_SOCKET_OPTION48        )49        self.assertTrue(tcp_nodelay_found)50 51    def test_recv_normal(self):52        """Test normal recv operation"""53        mock_sock = Mock()54        mock_sock.recv.return_value = b"test data"55 56        result = recv(mock_sock, 9)57 58        self.assertEqual(result, b"test data")59        mock_sock.recv.assert_called_once_with(9)60 61    def test_recv_timeout_error(self):62        """Test recv with TimeoutError"""63        mock_sock = Mock()64        mock_sock.recv.side_effect = TimeoutError("Connection timed out")65 66        with self.assertRaises(WebSocketTimeoutException) as cm:67            recv(mock_sock, 9)68 69        self.assertEqual(str(cm.exception), "Connection timed out")70 71    def test_recv_socket_timeout(self):72        """Test recv with socket.timeout"""73        mock_sock = Mock()74        timeout_exc = socket.timeout("Socket timed out")75        timeout_exc.args = ("Socket timed out",)76        mock_sock.recv.side_effect = timeout_exc77        mock_sock.gettimeout.return_value = 30.078 79        with self.assertRaises(WebSocketTimeoutException) as cm:80            recv(mock_sock, 9)81 82        # In Python 3.10+, socket.timeout is a subclass of TimeoutError83        # so it's caught by the TimeoutError handler with hardcoded message84        # In Python 3.9, socket.timeout is caught by socket.timeout handler85        # which preserves the original message86        import sys87 88        if sys.version_info >= (3, 10):89            self.assertEqual(str(cm.exception), "Connection timed out")90        else:91            self.assertEqual(str(cm.exception), "Socket timed out")92 93    def test_recv_ssl_timeout(self):94        """Test recv with SSL timeout error"""95        mock_sock = Mock()96        ssl_exc = SSLError("The operation timed out")97        ssl_exc.args = ("The operation timed out",)98        mock_sock.recv.side_effect = ssl_exc99 100        with self.assertRaises(WebSocketTimeoutException) as cm:101            recv(mock_sock, 9)102 103        self.assertEqual(str(cm.exception), "The operation timed out")104 105    def test_recv_ssl_non_timeout_error(self):106        """Test recv with SSL non-timeout error"""107        mock_sock = Mock()108        ssl_exc = SSLError("SSL certificate error")109        ssl_exc.args = ("SSL certificate error",)110        mock_sock.recv.side_effect = ssl_exc111 112        # Should re-raise the original SSL error113        with self.assertRaises(SSLError):114            recv(mock_sock, 9)115 116    def test_recv_empty_response(self):117        """Test recv with empty response (connection closed)"""118        mock_sock = Mock()119        mock_sock.recv.return_value = b""120 121        with self.assertRaises(WebSocketConnectionClosedException) as cm:122            recv(mock_sock, 9)123 124        self.assertEqual(str(cm.exception), "Connection to remote host was lost.")125 126    def test_recv_ssl_want_read_error(self):127        """Test recv with SSLWantReadError (should retry)"""128        mock_sock = Mock()129 130        # First call raises SSLWantReadError, second call succeeds131        mock_sock.recv.side_effect = [SSLWantReadError(), b"data after retry"]132 133        with patch("selectors.DefaultSelector") as mock_selector_class:134            mock_selector = Mock()135            mock_selector_class.return_value = mock_selector136            mock_selector.select.return_value = [True]  # Ready to read137 138            result = recv(mock_sock, 100)139 140            self.assertEqual(result, b"data after retry")141            mock_selector.register.assert_called()142            mock_selector.close.assert_called()143 144    def test_recv_ssl_want_read_timeout(self):145        """Test recv with SSLWantReadError that times out"""146        mock_sock = Mock()147        mock_sock.recv.side_effect = SSLWantReadError()148        mock_sock.gettimeout.return_value = 1.0149 150        with patch("selectors.DefaultSelector") as mock_selector_class:151            mock_selector = Mock()152            mock_selector_class.return_value = mock_selector153            mock_selector.select.return_value = []  # Timeout154 155            with self.assertRaises(WebSocketTimeoutException):156                recv(mock_sock, 100)157 158    def test_recv_line(self):159        """Test recv_line functionality"""160        mock_sock = Mock()161 162        # Mock recv to return one character at a time163        recv_calls = [b"H", b"e", b"l", b"l", b"o", b"\n"]164 165        with patch("websocket._socket.recv", side_effect=recv_calls) as mock_recv:166            result = recv_line(mock_sock)167 168            self.assertEqual(result, b"Hello\n")169            self.assertEqual(mock_recv.call_count, 6)170 171    def test_send_normal(self):172        """Test normal send operation"""173        mock_sock = Mock()174        mock_sock.send.return_value = 9175        mock_sock.gettimeout.return_value = 30.0176 177        result = send(mock_sock, b"test data")178 179        self.assertEqual(result, 9)180        mock_sock.send.assert_called_with(b"test data")181 182    def test_send_zero_timeout(self):183        """Test send with zero timeout (non-blocking)"""184        mock_sock = Mock()185        mock_sock.send.return_value = 9186        mock_sock.gettimeout.return_value = 0187 188        result = send(mock_sock, b"test data")189 190        self.assertEqual(result, 9)191        mock_sock.send.assert_called_once_with(b"test data")192 193    def test_send_ssl_eof_error(self):194        """Test send with SSLEOFError"""195        mock_sock = Mock()196        mock_sock.gettimeout.return_value = 30.0197        mock_sock.send.side_effect = SSLEOFError("Connection closed")198 199        with self.assertRaises(WebSocketConnectionClosedException) as cm:200            send(mock_sock, b"test data")201 202        self.assertEqual(str(cm.exception), "socket is already closed.")203 204    def test_send_ssl_want_write_error(self):205        """Test send with SSLWantWriteError (should retry)"""206        mock_sock = Mock()207        mock_sock.gettimeout.return_value = 30.0208 209        # First call raises SSLWantWriteError, second call succeeds210        mock_sock.send.side_effect = [SSLWantWriteError(), 9]211 212        with patch("selectors.DefaultSelector") as mock_selector_class:213            mock_selector = Mock()214            mock_selector_class.return_value = mock_selector215            mock_selector.select.return_value = [True]  # Ready to write216 217            result = send(mock_sock, b"test data")218 219            self.assertEqual(result, 9)220            mock_selector.register.assert_called()221            mock_selector.close.assert_called()222 223    def test_send_socket_eagain_error(self):224        """Test send with EAGAIN error (should retry)"""225        mock_sock = Mock()226        mock_sock.gettimeout.return_value = 30.0227 228        # Create socket error with EAGAIN229        eagain_error = socket.error("Resource temporarily unavailable")230        eagain_error.errno = errno.EAGAIN231        eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")232 233        # First call raises EAGAIN, second call succeeds234        mock_sock.send.side_effect = [eagain_error, 9]235 236        with patch("selectors.DefaultSelector") as mock_selector_class:237            mock_selector = Mock()238            mock_selector_class.return_value = mock_selector239            mock_selector.select.return_value = [True]  # Ready to write240 241            result = send(mock_sock, b"test data")242 243            self.assertEqual(result, 9)244 245    def test_send_socket_ewouldblock_error(self):246        """Test send with EWOULDBLOCK error (should retry)"""247        mock_sock = Mock()248        mock_sock.gettimeout.return_value = 30.0249 250        # Create socket error with EWOULDBLOCK251        ewouldblock_error = socket.error("Operation would block")252        ewouldblock_error.errno = errno.EWOULDBLOCK253        ewouldblock_error.args = (errno.EWOULDBLOCK, "Operation would block")254 255        # First call raises EWOULDBLOCK, second call succeeds256        mock_sock.send.side_effect = [ewouldblock_error, 9]257 258        with patch("selectors.DefaultSelector") as mock_selector_class:259            mock_selector = Mock()260            mock_selector_class.return_value = mock_selector261            mock_selector.select.return_value = [True]  # Ready to write262 263            result = send(mock_sock, b"test data")264 265            self.assertEqual(result, 9)266 267    def test_send_socket_other_error(self):268        """Test send with other socket error (should raise)"""269        mock_sock = Mock()270        mock_sock.gettimeout.return_value = 30.0271 272        # Create socket error with different errno273        other_error = socket.error("Connection reset by peer")274        other_error.errno = errno.ECONNRESET275        other_error.args = (errno.ECONNRESET, "Connection reset by peer")276 277        mock_sock.send.side_effect = other_error278 279        with self.assertRaises(socket.error):280            send(mock_sock, b"test data")281 282    def test_send_socket_error_no_errno(self):283        """Test send with socket error that has no errno"""284        mock_sock = Mock()285        mock_sock.gettimeout.return_value = 30.0286 287        # Create socket error without errno attribute288        no_errno_error = socket.error("Generic socket error")289        no_errno_error.args = ("Generic socket error",)290 291        mock_sock.send.side_effect = no_errno_error292 293        with self.assertRaises(socket.error):294            send(mock_sock, b"test data")295 296    def test_send_write_timeout(self):297        """Test send write operation timeout"""298        mock_sock = Mock()299        mock_sock.gettimeout.return_value = 30.0300 301        # First call raises EAGAIN302        eagain_error = socket.error("Resource temporarily unavailable")303        eagain_error.errno = errno.EAGAIN304        eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")305 306        mock_sock.send.side_effect = eagain_error307 308        with patch("selectors.DefaultSelector") as mock_selector_class:309            mock_selector = Mock()310            mock_selector_class.return_value = mock_selector311            mock_selector.select.return_value = []  # Timeout - nothing ready312 313            result = send(mock_sock, b"test data")314 315            # Should return 0 when write times out316            self.assertEqual(result, 0)317 318    def test_send_string_data(self):319        """Test send with string data (should be encoded)"""320        mock_sock = Mock()321        mock_sock.send.return_value = 9322        mock_sock.gettimeout.return_value = 30.0323 324        result = send(mock_sock, "test data")325 326        self.assertEqual(result, 9)327        mock_sock.send.assert_called_with(b"test data")328 329    def test_send_partial_send_retry(self):330        """Test send retry mechanism"""331        mock_sock = Mock()332        mock_sock.gettimeout.return_value = 30.0333 334        # Create a scenario where send succeeds after selector retry335        eagain_error = socket.error("Resource temporarily unavailable")336        eagain_error.errno = errno.EAGAIN337        eagain_error.args = (errno.EAGAIN, "Resource temporarily unavailable")338 339        # Mock the internal _send function behavior340        mock_sock.send.side_effect = [eagain_error, 9]341 342        with patch("selectors.DefaultSelector") as mock_selector_class:343            mock_selector = Mock()344            mock_selector_class.return_value = mock_selector345            mock_selector.select.return_value = [True]  # Socket ready for writing346 347            result = send(mock_sock, b"test data")348 349            self.assertEqual(result, 9)350            # Verify selector was used for retry mechanism351            mock_selector.register.assert_called()352            mock_selector.select.assert_called()353            mock_selector.close.assert_called()354 355 356if __name__ == "__main__":357    unittest.main()358