test_dispatcher.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385
  1. # -*- coding: utf-8 -*-
  2. import socket
  3. import unittest
  4. from unittest.mock import Mock, patch, MagicMock
  5. import threading
  6. import time
  7. import websocket
  8. from websocket._dispatcher import (
  9. Dispatcher,
  10. DispatcherBase,
  11. SSLDispatcher,
  12. WrappedDispatcher,
  13. )
  14. """
  15. test_dispatcher.py
  16. websocket - WebSocket client library for Python
  17. Copyright 2025 engn33r
  18. Licensed under the Apache License, Version 2.0 (the "License");
  19. you may not use this file except in compliance with the License.
  20. You may obtain a copy of the License at
  21. http://www.apache.org/licenses/LICENSE-2.0
  22. Unless required by applicable law or agreed to in writing, software
  23. distributed under the License is distributed on an "AS IS" BASIS,
  24. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  25. See the License for the specific language governing permissions and
  26. limitations under the License.
  27. """
  28. class MockApp:
  29. """Mock WebSocketApp for testing"""
  30. def __init__(self):
  31. self.keep_running = True
  32. self.sock = Mock()
  33. self.sock.sock = Mock()
  34. class MockSocket:
  35. """Mock socket for testing"""
  36. def __init__(self):
  37. self.pending_return = False
  38. def pending(self):
  39. return self.pending_return
  40. class MockDispatcher:
  41. """Mock external dispatcher for WrappedDispatcher testing"""
  42. def __init__(self):
  43. self.signal_calls = []
  44. self.abort_calls = []
  45. self.read_calls = []
  46. self.buffwrite_calls = []
  47. self.timeout_calls = []
  48. def signal(self, sig, handler):
  49. self.signal_calls.append((sig, handler))
  50. def abort(self):
  51. self.abort_calls.append(True)
  52. def read(self, sock, callback):
  53. self.read_calls.append((sock, callback))
  54. def buffwrite(self, sock, data, send_func, disconnect_handler):
  55. self.buffwrite_calls.append((sock, data, send_func, disconnect_handler))
  56. def timeout(self, seconds, callback, *args):
  57. self.timeout_calls.append((seconds, callback, args))
  58. class DispatcherTest(unittest.TestCase):
  59. def setUp(self):
  60. self.app = MockApp()
  61. def test_dispatcher_base_init(self):
  62. """Test DispatcherBase initialization"""
  63. dispatcher = DispatcherBase(self.app, 30.0)
  64. self.assertEqual(dispatcher.app, self.app)
  65. self.assertEqual(dispatcher.ping_timeout, 30.0)
  66. def test_dispatcher_base_timeout(self):
  67. """Test DispatcherBase timeout method"""
  68. dispatcher = DispatcherBase(self.app, 30.0)
  69. callback = Mock()
  70. # Test with seconds=None (should call callback immediately)
  71. dispatcher.timeout(None, callback)
  72. callback.assert_called_once()
  73. # Test with seconds > 0 (would sleep in real implementation)
  74. callback.reset_mock()
  75. start_time = time.time()
  76. dispatcher.timeout(0.1, callback)
  77. elapsed = time.time() - start_time
  78. callback.assert_called_once()
  79. self.assertGreaterEqual(elapsed, 0.05) # Allow some tolerance
  80. def test_dispatcher_base_reconnect(self):
  81. """Test DispatcherBase reconnect method"""
  82. dispatcher = DispatcherBase(self.app, 30.0)
  83. reconnector = Mock()
  84. # Test normal reconnect
  85. dispatcher.reconnect(1, reconnector)
  86. reconnector.assert_called_once_with(reconnecting=True)
  87. # Test reconnect with KeyboardInterrupt
  88. reconnector.reset_mock()
  89. reconnector.side_effect = KeyboardInterrupt("User interrupted")
  90. with self.assertRaises(KeyboardInterrupt):
  91. dispatcher.reconnect(1, reconnector)
  92. def test_dispatcher_base_send(self):
  93. """Test DispatcherBase send method"""
  94. dispatcher = DispatcherBase(self.app, 30.0)
  95. mock_sock = Mock()
  96. test_data = b"test data"
  97. with patch("websocket._dispatcher.send") as mock_send:
  98. mock_send.return_value = len(test_data)
  99. result = dispatcher.send(mock_sock, test_data)
  100. mock_send.assert_called_once_with(mock_sock, test_data)
  101. self.assertEqual(result, len(test_data))
  102. def test_dispatcher_read(self):
  103. """Test Dispatcher read method"""
  104. dispatcher = Dispatcher(self.app, 5.0)
  105. read_callback = Mock(return_value=True)
  106. check_callback = Mock()
  107. mock_sock = Mock()
  108. # Mock the selector to control the loop
  109. with patch("selectors.DefaultSelector") as mock_selector_class:
  110. mock_selector = Mock()
  111. mock_selector_class.return_value = mock_selector
  112. # Make select return immediately (timeout)
  113. mock_selector.select.return_value = []
  114. # Stop after first iteration
  115. def side_effect(*args):
  116. self.app.keep_running = False
  117. return []
  118. mock_selector.select.side_effect = side_effect
  119. dispatcher.read(mock_sock, read_callback, check_callback)
  120. # Verify selector was used correctly
  121. mock_selector.register.assert_called()
  122. mock_selector.select.assert_called_with(5.0)
  123. mock_selector.close.assert_called()
  124. check_callback.assert_called()
  125. def test_dispatcher_read_with_data(self):
  126. """Test Dispatcher read method when data is available"""
  127. dispatcher = Dispatcher(self.app, 5.0)
  128. read_callback = Mock(return_value=True)
  129. check_callback = Mock()
  130. mock_sock = Mock()
  131. with patch("selectors.DefaultSelector") as mock_selector_class:
  132. mock_selector = Mock()
  133. mock_selector_class.return_value = mock_selector
  134. # First call returns data, second call stops the loop
  135. call_count = 0
  136. def select_side_effect(*args):
  137. nonlocal call_count
  138. call_count += 1
  139. if call_count == 1:
  140. return [True] # Data available
  141. else:
  142. self.app.keep_running = False
  143. return []
  144. mock_selector.select.side_effect = select_side_effect
  145. dispatcher.read(mock_sock, read_callback, check_callback)
  146. read_callback.assert_called()
  147. check_callback.assert_called()
  148. def test_ssl_dispatcher_read(self):
  149. """Test SSLDispatcher read method"""
  150. dispatcher = SSLDispatcher(self.app, 5.0)
  151. read_callback = Mock(return_value=True)
  152. check_callback = Mock()
  153. # Mock socket with pending data
  154. mock_ssl_sock = MockSocket()
  155. self.app.sock.sock = mock_ssl_sock
  156. with patch("selectors.DefaultSelector") as mock_selector_class:
  157. mock_selector = Mock()
  158. mock_selector_class.return_value = mock_selector
  159. mock_selector.select.return_value = []
  160. # Stop after first iteration
  161. def side_effect(*args):
  162. self.app.keep_running = False
  163. return []
  164. mock_selector.select.side_effect = side_effect
  165. dispatcher.read(None, read_callback, check_callback)
  166. mock_selector.register.assert_called()
  167. check_callback.assert_called()
  168. def test_ssl_dispatcher_select_with_pending(self):
  169. """Test SSLDispatcher select method with pending data"""
  170. dispatcher = SSLDispatcher(self.app, 5.0)
  171. mock_ssl_sock = MockSocket()
  172. mock_ssl_sock.pending_return = True
  173. self.app.sock.sock = mock_ssl_sock
  174. mock_selector = Mock()
  175. result = dispatcher.select(None, mock_selector)
  176. # When pending() returns True, should return [sock]
  177. self.assertEqual(result, [mock_ssl_sock])
  178. def test_ssl_dispatcher_select_without_pending(self):
  179. """Test SSLDispatcher select method without pending data"""
  180. dispatcher = SSLDispatcher(self.app, 5.0)
  181. mock_ssl_sock = MockSocket()
  182. mock_ssl_sock.pending_return = False
  183. self.app.sock.sock = mock_ssl_sock
  184. mock_selector = Mock()
  185. mock_selector.select.return_value = [(mock_ssl_sock, None)]
  186. result = dispatcher.select(None, mock_selector)
  187. # Should return the first element of first result tuple
  188. self.assertEqual(result, mock_ssl_sock)
  189. mock_selector.select.assert_called_with(5.0)
  190. def test_ssl_dispatcher_select_no_results(self):
  191. """Test SSLDispatcher select method with no results"""
  192. dispatcher = SSLDispatcher(self.app, 5.0)
  193. mock_ssl_sock = MockSocket()
  194. mock_ssl_sock.pending_return = False
  195. self.app.sock.sock = mock_ssl_sock
  196. mock_selector = Mock()
  197. mock_selector.select.return_value = []
  198. result = dispatcher.select(None, mock_selector)
  199. # Should return None when no results (function doesn't return anything when len(r) == 0)
  200. self.assertIsNone(result)
  201. def test_wrapped_dispatcher_init(self):
  202. """Test WrappedDispatcher initialization"""
  203. mock_dispatcher = MockDispatcher()
  204. handle_disconnect = Mock()
  205. wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
  206. self.assertEqual(wrapped.app, self.app)
  207. self.assertEqual(wrapped.ping_timeout, 10.0)
  208. self.assertEqual(wrapped.dispatcher, mock_dispatcher)
  209. self.assertEqual(wrapped.handleDisconnect, handle_disconnect)
  210. # Should have set up signal handler
  211. self.assertEqual(len(mock_dispatcher.signal_calls), 1)
  212. sig, handler = mock_dispatcher.signal_calls[0]
  213. self.assertEqual(sig, 2) # SIGINT
  214. self.assertEqual(handler, mock_dispatcher.abort)
  215. def test_wrapped_dispatcher_read(self):
  216. """Test WrappedDispatcher read method"""
  217. mock_dispatcher = MockDispatcher()
  218. handle_disconnect = Mock()
  219. wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
  220. mock_sock = Mock()
  221. read_callback = Mock()
  222. check_callback = Mock()
  223. wrapped.read(mock_sock, read_callback, check_callback)
  224. # Should delegate to wrapped dispatcher
  225. self.assertEqual(len(mock_dispatcher.read_calls), 1)
  226. self.assertEqual(mock_dispatcher.read_calls[0], (mock_sock, read_callback))
  227. # Should call timeout for ping_timeout
  228. self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
  229. timeout_call = mock_dispatcher.timeout_calls[0]
  230. self.assertEqual(timeout_call[0], 10.0) # timeout seconds
  231. self.assertEqual(timeout_call[1], check_callback) # callback
  232. def test_wrapped_dispatcher_read_no_ping_timeout(self):
  233. """Test WrappedDispatcher read method without ping timeout"""
  234. mock_dispatcher = MockDispatcher()
  235. handle_disconnect = Mock()
  236. wrapped = WrappedDispatcher(self.app, None, mock_dispatcher, handle_disconnect)
  237. mock_sock = Mock()
  238. read_callback = Mock()
  239. check_callback = Mock()
  240. wrapped.read(mock_sock, read_callback, check_callback)
  241. # Should delegate to wrapped dispatcher
  242. self.assertEqual(len(mock_dispatcher.read_calls), 1)
  243. # Should NOT call timeout when ping_timeout is None
  244. self.assertEqual(len(mock_dispatcher.timeout_calls), 0)
  245. def test_wrapped_dispatcher_send(self):
  246. """Test WrappedDispatcher send method"""
  247. mock_dispatcher = MockDispatcher()
  248. handle_disconnect = Mock()
  249. wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
  250. mock_sock = Mock()
  251. test_data = b"test data"
  252. with patch("websocket._dispatcher.send") as mock_send:
  253. result = wrapped.send(mock_sock, test_data)
  254. # Should delegate to dispatcher.buffwrite
  255. self.assertEqual(len(mock_dispatcher.buffwrite_calls), 1)
  256. call = mock_dispatcher.buffwrite_calls[0]
  257. self.assertEqual(call[0], mock_sock)
  258. self.assertEqual(call[1], test_data)
  259. self.assertEqual(call[2], mock_send)
  260. self.assertEqual(call[3], handle_disconnect)
  261. # Should return data length
  262. self.assertEqual(result, len(test_data))
  263. def test_wrapped_dispatcher_timeout(self):
  264. """Test WrappedDispatcher timeout method"""
  265. mock_dispatcher = MockDispatcher()
  266. handle_disconnect = Mock()
  267. wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
  268. callback = Mock()
  269. args = ("arg1", "arg2")
  270. wrapped.timeout(5.0, callback, *args)
  271. # Should delegate to wrapped dispatcher
  272. self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
  273. call = mock_dispatcher.timeout_calls[0]
  274. self.assertEqual(call[0], 5.0)
  275. self.assertEqual(call[1], callback)
  276. self.assertEqual(call[2], args)
  277. def test_wrapped_dispatcher_reconnect(self):
  278. """Test WrappedDispatcher reconnect method"""
  279. mock_dispatcher = MockDispatcher()
  280. handle_disconnect = Mock()
  281. wrapped = WrappedDispatcher(self.app, 10.0, mock_dispatcher, handle_disconnect)
  282. reconnector = Mock()
  283. wrapped.reconnect(3, reconnector)
  284. # Should delegate to timeout method with reconnect=True
  285. self.assertEqual(len(mock_dispatcher.timeout_calls), 1)
  286. call = mock_dispatcher.timeout_calls[0]
  287. self.assertEqual(call[0], 3)
  288. self.assertEqual(call[1], reconnector)
  289. self.assertEqual(call[2], (True,))
  290. if __name__ == "__main__":
  291. unittest.main()