codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2import socket3import unittest4from unittest.mock import Mock, patch, MagicMock5import threading6import time7 8import websocket9from websocket._dispatcher import (10    Dispatcher,11    DispatcherBase,12    SSLDispatcher,13    WrappedDispatcher,14)15 16"""17test_dispatcher.py18websocket - WebSocket client library for Python19 20Copyright 2025 engn33r21 22Licensed under the Apache License, Version 2.0 (the "License");23you may not use this file except in compliance with the License.24You may obtain a copy of the License at25 26    http://www.apache.org/licenses/LICENSE-2.027 28Unless required by applicable law or agreed to in writing, software29distributed under the License is distributed on an "AS IS" BASIS,30WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.31See the License for the specific language governing permissions and32limitations under the License.33"""34 35class MockApp:36    """Mock WebSocketApp for testing"""37 38    def __init__(self):39        self.keep_running = True40        self.sock = Mock()41        self.sock.sock = Mock()42 43 44class MockSocket:45    """Mock socket for testing"""46 47    def __init__(self):48        self.pending_return = False49 50    def pending(self):51        return self.pending_return52 53 54class MockDispatcher:55    """Mock external dispatcher for WrappedDispatcher testing"""56 57    def __init__(self):58        self.signal_calls = []59        self.abort_calls = []60        self.read_calls = []61        self.buffwrite_calls = []62        self.timeout_calls = []63 64    def signal(self, sig, handler):65        self.signal_calls.append((sig, handler))66 67    def abort(self):68        self.abort_calls.append(True)69 70    def read(self, sock, callback):71        self.read_calls.append((sock, callback))72 73    def buffwrite(self, sock, data, send_func, disconnect_handler):74        self.buffwrite_calls.append((sock, data, send_func, disconnect_handler))75 76    def timeout(self, seconds, callback, *args):77        self.timeout_calls.append((seconds, callback, args))78 79 80class DispatcherTest(unittest.TestCase):81    def setUp(self):82        self.app = MockApp()83 84    def test_dispatcher_base_init(self):85        """Test DispatcherBase initialization"""86        dispatcher = DispatcherBase(self.app, 30.0)87 88        self.assertEqual(dispatcher.app, self.app)89        self.assertEqual(dispatcher.ping_timeout, 30.0)90 91    def test_dispatcher_base_timeout(self):92        """Test DispatcherBase timeout method"""93        dispatcher = DispatcherBase(self.app, 30.0)94        callback = Mock()95 96        # Test with seconds=None (should call callback immediately)97        dispatcher.timeout(None, callback)98        callback.assert_called_once()99 100        # Test with seconds > 0 (would sleep in real implementation)101        callback.reset_mock()102        start_time = time.time()103        dispatcher.timeout(0.1, callback)104        elapsed = time.time() - start_time105 106        callback.assert_called_once()107        self.assertGreaterEqual(elapsed, 0.05)  # Allow some tolerance108 109    def test_dispatcher_base_reconnect(self):110        """Test DispatcherBase reconnect method"""111        dispatcher = DispatcherBase(self.app, 30.0)112        reconnector = Mock()113 114        # Test normal reconnect115        dispatcher.reconnect(1, reconnector)116        reconnector.assert_called_once_with(reconnecting=True)117 118        # Test reconnect with KeyboardInterrupt119        reconnector.reset_mock()120        reconnector.side_effect = KeyboardInterrupt("User interrupted")121 122        with self.assertRaises(KeyboardInterrupt):123            dispatcher.reconnect(1, reconnector)124 125    def test_dispatcher_base_send(self):126        """Test DispatcherBase send method"""127        dispatcher = DispatcherBase(self.app, 30.0)128        mock_sock = Mock()129        test_data = b"test data"130 131        with patch("websocket._dispatcher.send") as mock_send:132            mock_send.return_value = len(test_data)133            result = dispatcher.send(mock_sock, test_data)134 135            mock_send.assert_called_once_with(mock_sock, test_data)136            self.assertEqual(result, len(test_data))137 138    def test_dispatcher_read(self):139        """Test Dispatcher read method"""140        dispatcher = Dispatcher(self.app, 5.0)141        read_callback = Mock(return_value=True)142        check_callback = Mock()143        mock_sock = Mock()144 145        # Mock the selector to control the loop146        with patch("selectors.DefaultSelector") as mock_selector_class:147            mock_selector = Mock()148            mock_selector_class.return_value = mock_selector149 150            # Make select return immediately (timeout)151            mock_selector.select.return_value = []152 153            # Stop after first iteration154            def side_effect(*args):155                self.app.keep_running = False156                return []157 158            mock_selector.select.side_effect = side_effect159 160            dispatcher.read(mock_sock, read_callback, check_callback)161 162            # Verify selector was used correctly163            mock_selector.register.assert_called()164            mock_selector.select.assert_called_with(5.0)165            mock_selector.close.assert_called()166            check_callback.assert_called()167 168    def test_dispatcher_read_with_data(self):169        """Test Dispatcher read method when data is available"""170        dispatcher = Dispatcher(self.app, 5.0)171        read_callback = Mock(return_value=True)172        check_callback = Mock()173        mock_sock = Mock()174 175        with patch("selectors.DefaultSelector") as mock_selector_class:176            mock_selector = Mock()177            mock_selector_class.return_value = mock_selector178 179            # First call returns data, second call stops the loop180            call_count = 0181 182            def select_side_effect(*args):183                nonlocal call_count184                call_count += 1185                if call_count == 1:186                    return [True]  # Data available187                else:188                    self.app.keep_running = False189                    return []190 191            mock_selector.select.side_effect = select_side_effect192 193            dispatcher.read(mock_sock, read_callback, check_callback)194 195            read_callback.assert_called()196            check_callback.assert_called()197 198    def test_ssl_dispatcher_read(self):199        """Test SSLDispatcher read method"""200        dispatcher = SSLDispatcher(self.app, 5.0)201        read_callback = Mock(return_value=True)202        check_callback = Mock()203 204        # Mock socket with pending data205        mock_ssl_sock = MockSocket()206        self.app.sock.sock = mock_ssl_sock207 208        with patch("selectors.DefaultSelector") as mock_selector_class:209            mock_selector = Mock()210            mock_selector_class.return_value = mock_selector211            mock_selector.select.return_value = []212 213            # Stop after first iteration214            def side_effect(*args):215                self.app.keep_running = False216                return []217 218            mock_selector.select.side_effect = side_effect219 220            dispatcher.read(None, read_callback, check_callback)221 222            mock_selector.register.assert_called()223            check_callback.assert_called()224 225    def test_ssl_dispatcher_select_with_pending(self):226        """Test SSLDispatcher select method with pending data"""227        dispatcher = SSLDispatcher(self.app, 5.0)228        mock_ssl_sock = MockSocket()229        mock_ssl_sock.pending_return = True230        self.app.sock.sock = mock_ssl_sock231        mock_selector = Mock()232 233        result = dispatcher.select(None, mock_selector)234 235        # When pending() returns True, should return [sock]236        self.assertEqual(result, [mock_ssl_sock])237 238    def test_ssl_dispatcher_select_without_pending(self):239        """Test SSLDispatcher select method without pending data"""240        dispatcher = SSLDispatcher(self.app, 5.0)241        mock_ssl_sock = MockSocket()242        mock_ssl_sock.pending_return = False243        self.app.sock.sock = mock_ssl_sock244        mock_selector = Mock()245        mock_selector.select.return_value = [(mock_ssl_sock, None)]246 247        result = dispatcher.select(None, mock_selector)248 249        # Should return the first element of first result tuple250        self.assertEqual(result, mock_ssl_sock)251        mock_selector.select.assert_called_with(5.0)252 253    def test_ssl_dispatcher_select_no_results(self):254        """Test SSLDispatcher select method with no results"""255        dispatcher = SSLDispatcher(self.app, 5.0)256        mock_ssl_sock = MockSocket()257        mock_ssl_sock.pending_return = False258        self.app.sock.sock = mock_ssl_sock259        mock_selector = Mock()260        mock_selector.select.return_value = []261 262        result = dispatcher.select(None, mock_selector)263 264        # Should return None when no results (function doesn't return anything when len(r) == 0)265        self.assertIsNone(result)266 267    def test_wrapped_dispatcher_init(self):268        """Test WrappedDispatcher initialization"""269        mock_dispatcher = MockDispatcher()270        handle_disconnect = Mock()271 272        wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)273 274        self.assertEqual(wrapped.app, self.app)275        self.assertEqual(wrapped.ping_timeout, 10.0)276        self.assertEqual(wrapped.dispatcher, mock_dispatcher)277        self.assertEqual(wrapped.handleDisconnect, handle_disconnect)278 279        # Should have set up signal handler280        self.assertEqual(len(mock_dispatcher.signal_calls), 1)281        sig, handler = mock_dispatcher.signal_calls[0]282        self.assertEqual(sig, 2)  # SIGINT283        self.assertEqual(handler, mock_dispatcher.abort)284 285    def test_wrapped_dispatcher_read(self):286        """Test WrappedDispatcher read method"""287        mock_dispatcher = MockDispatcher()288        handle_disconnect = Mock()289        wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)290 291        mock_sock = Mock()292        read_callback = Mock()293        check_callback = Mock()294 295        wrapped.read(mock_sock, read_callback, check_callback)296 297        # Should delegate to wrapped dispatcher298        self.assertEqual(len(mock_dispatcher.read_calls), 1)299        self.assertEqual(mock_dispatcher.read_calls[0], (mock_sock, read_callback))300 301        # Should call timeout for ping_timeout302        self.assertEqual(len(mock_dispatcher.timeout_calls), 1)303        timeout_call = mock_dispatcher.timeout_calls[0]304        self.assertEqual(timeout_call[0], 10.0)  # timeout seconds305        self.assertEqual(timeout_call[1], check_callback)  # callback306 307    def test_wrapped_dispatcher_read_no_ping_timeout(self):308        """Test WrappedDispatcher read method without ping timeout"""309        mock_dispatcher = MockDispatcher()310        handle_disconnect = Mock()311        wrapped = WrappedDispatcher(self.app, None, mock_dispatcher, handle_disconnect)312 313        mock_sock = Mock()314        read_callback = Mock()315        check_callback = Mock()316 317        wrapped.read(mock_sock, read_callback, check_callback)318 319        # Should delegate to wrapped dispatcher320        self.assertEqual(len(mock_dispatcher.read_calls), 1)321 322        # Should NOT call timeout when ping_timeout is None323        self.assertEqual(len(mock_dispatcher.timeout_calls), 0)324 325    def test_wrapped_dispatcher_send(self):326        """Test WrappedDispatcher send method"""327        mock_dispatcher = MockDispatcher()328        handle_disconnect = Mock()329        wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)330 331        mock_sock = Mock()332        test_data = b"test data"333 334        with patch("websocket._dispatcher.send") as mock_send:335            result = wrapped.send(mock_sock, test_data)336 337            # Should delegate to dispatcher.buffwrite338            self.assertEqual(len(mock_dispatcher.buffwrite_calls), 1)339            call = mock_dispatcher.buffwrite_calls[0]340            self.assertEqual(call[0], mock_sock)341            self.assertEqual(call[1], test_data)342            self.assertEqual(call[2], mock_send)343            self.assertEqual(call[3], handle_disconnect)344 345            # Should return data length346            self.assertEqual(result, len(test_data))347 348    def test_wrapped_dispatcher_timeout(self):349        """Test WrappedDispatcher timeout method"""350        mock_dispatcher = MockDispatcher()351        handle_disconnect = Mock()352        wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)353 354        callback = Mock()355        args = ("arg1", "arg2")356 357        wrapped.timeout(5.0, callback, *args)358 359        # Should delegate to wrapped dispatcher360        self.assertEqual(len(mock_dispatcher.timeout_calls), 1)361        call = mock_dispatcher.timeout_calls[0]362        self.assertEqual(call[0], 5.0)363        self.assertEqual(call[1], callback)364        self.assertEqual(call[2], args)365 366    def test_wrapped_dispatcher_reconnect(self):367        """Test WrappedDispatcher reconnect method"""368        mock_dispatcher = MockDispatcher()369        handle_disconnect = Mock()370        wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)371 372        reconnector = Mock()373 374        wrapped.reconnect(3, reconnector)375 376        # Should delegate to timeout method with reconnect=True377        self.assertEqual(len(mock_dispatcher.timeout_calls), 1)378        call = mock_dispatcher.timeout_calls[0]379        self.assertEqual(call[0], 3)380        self.assertEqual(call[1], reconnector)381        self.assertEqual(call[2], (True,))382 383 384if __name__ == "__main__":385    unittest.main()386