diff --git a/test/test_websocket.py b/test/test_websocket.py index 96fb384..ab319bf 100644 --- a/test/test_websocket.py +++ b/test/test_websocket.py @@ -178,6 +178,28 @@ def test_sending_ping(self): m.sendall.assert_called_once_with(tm) + @patch("ws4py.websocket.Heartbeat") + def test_run(self, mocker): + mocked_sock = MagicMock() + mocked_opened = MagicMock() + mocked_once = MagicMock(return_value=False) # False to break the loop + mocked_terminate = MagicMock() + + ws = WebSocket(sock=mocked_sock) + assert ws.sock_timeout is None + + with patch.multiple(ws, + opened=mocked_opened, + once=mocked_once, + terminate=mocked_terminate, + stream=MagicMock(), + ): + ws.run() + mocked_sock.settimeout.assert_called_with(None) + mocked_opened.assert_called() + mocked_once.assert_called() + mocked_terminate.assert_called() + if __name__ == '__main__': suite = unittest.TestSuite() loader = unittest.TestLoader() diff --git a/ws4py/websocket.py b/ws4py/websocket.py index 61f8c33..7cdc5e6 100644 --- a/ws4py/websocket.py +++ b/ws4py/websocket.py @@ -141,6 +141,17 @@ def __init__(self, sock, protocols=None, extensions=None, environ=None, heartbea "Internal buffer to get around SSL problems" self.buf = b'' + self.sock_timeout = None + """ + Used to set socket.settimeout(value): + From: https://docs.python.org/3.11/library/socket.html#socket.socket.settimeout + The value argument can be a nonnegative floating point number expressing seconds, or None. + If a non-zero value is given, subsequent socket operations will raise a timeout exception + if the timeout period value has elapsed before the operation has completed. + If zero is given, the socket is put in non-blocking mode. + If None is given, the socket is put in blocking mode. + """ + self._local_address = None self._peer_address = None @@ -515,10 +526,12 @@ def run(self): we initiate the closing of the connection with the appropiate error code. - This method is blocking and should likely be run - in a thread. + The self.sock_timeout determines whether this method + is blocking, or can timeout on reads. If a timeout + occurs, the unhandled_error function will be called + It should likely be run in a thread. """ - self.sock.setblocking(True) + self.sock.settimeout(self.sock_timeout) with Heartbeat(self, frequency=self.heartbeat_freq): s = self.stream