ws client hardening

Author: tmunzer-AIDE

src/mistapi/websockets/__ws_client.py

@@ -70,6 +70,7 @@ def __init__(
         )  # tracks whether the WebSocket connection is currently open
         self._user_disconnect = threading.Event()
         self._finished = threading.Event()
+        self._finished.set()  # not running initially
         self._reconnect_attempts = 0
         self._last_close_code: int | None = None
         self._last_close_msg: str | None = None
@@ -164,7 +165,6 @@ def _handle_open(self, ws: websocket.WebSocketApp) -> None:
                 ws.send(json.dumps({"subscribe": channel}))
         except Exception as exc:
             logger.error("Subscription send failed: %s", exc)
-            self._handle_error(ws, exc)
             ws.close()
             return
         self._reconnect_attempts = 0
@@ -235,12 +235,10 @@ def connect(self, run_in_background: bool = True) -> None:
             If False, blocks the calling thread until disconnected.
         """
         with self._lock:
-            if self._connected.is_set() or (
-                self._thread is not None and self._thread.is_alive()
-            ):
+            if self._connected.is_set() or not self._finished.is_set():
                 raise RuntimeError("Already connected; call disconnect() first")
-            self._user_disconnect.clear()
             self._finished.clear()
+            self._user_disconnect.clear()
             self._reconnect_attempts = 0
             # Drain stale sentinel from previous connection
             while not self._queue.empty():
@@ -309,15 +307,14 @@ def _run_forever_safe(self) -> None:
                     except Exception:
                         pass
 
-            # Final close: put sentinel, call callback
-            self._queue.put(None)
+        finally:
+            self._queue.put(None)  # sentinel — unblocks receive()
+            self._finished.set()  # mark as not running — unblocks connect()
             if self._on_close_cb:
                 try:
                     self._on_close_cb(self._last_close_code, self._last_close_msg)
                 except Exception:
                     logger.exception("on_close callback raised")
-        finally:
-            self._finished.set()
 
     def disconnect(self, wait: bool = False, timeout: float | None = None) -> None:
         """Close the WebSocket connection.
@@ -336,7 +333,8 @@ def disconnect(self, wait: bool = False, timeout: float | None = None) -> None:
         if ws:
             ws.close()
         if wait and self._thread is not None:
-            self._thread.join(timeout=timeout)
+            if self._thread is not threading.current_thread():
+                self._thread.join(timeout=timeout)
 
     def receive(self) -> Generator[dict, None, None]:
         """
@@ -346,7 +344,13 @@ def receive(self) -> Generator[dict, None, None]:
         the server closes the connection).
 
         Intended for use after connect(run_in_background=True).
+        Cannot be used when an on_message callback is registered.
         """
+        if self._on_message_cb is not None:
+            raise RuntimeError(
+                "receive() cannot be used when an on_message callback is "
+                "registered; use one or the other"
+            )
         if self._auto_reconnect:
             while (
                 not self._connected.is_set()
@@ -357,11 +361,15 @@ def receive(self) -> Generator[dict, None, None]:
             if not self._connected.is_set():
                 return
         elif not self._connected.wait(timeout=10):
-            return
+            if not self._finished.is_set():
+                return
+            # Thread already finished — fall through to drain queued messages
         while True:
             try:
                 item = self._queue.get(timeout=1)
             except queue.Empty:
+                if self._finished.is_set() and self._queue.empty():
+                    break
                 if not self._connected.is_set() and self._queue.empty():
                     if (
                         self._auto_reconnect

tests/unit/test_websocket_client.py

@@ -1026,21 +1026,40 @@ def bad_cb(code, msg):
         # Sentinel should still be in the queue
         assert ws_client._queue.get_nowait() is None
 
+    @patch("mistapi.websockets.__ws_client.websocket.WebSocketApp")
+    def test_connect_from_on_close_callback(self, mock_ws_cls, mock_session) -> None:
+        """connect() must work when called from inside the on_close callback."""
+        mock_ws_cls.return_value = Mock()
+        client = _MistWebsocket(mock_session, channels=["/ch"])
+        errors = []
+
+        def on_close_cb(code, msg):
+            try:
+                client.connect(run_in_background=True)
+            except RuntimeError as exc:
+                errors.append(exc)
+            finally:
+                client.disconnect()
+
+        client.on_close(on_close_cb)
+        mock_ws = Mock()
+        client._ws = mock_ws
+        client._finished.clear()  # simulate connect() was called
+        client._run_forever_safe()
+
+        assert not errors, f"on_close callback raised: {errors}"
+
     def test_on_open_send_failure_closes_connection(self, ws_client) -> None:
         """If ws.send() raises during subscription, connection is closed."""
         mock_ws = Mock()
         mock_ws.send.side_effect = ConnectionError("send failed")
-        error_cb = Mock()
-        ws_client.on_error(error_cb)
 
         ws_client._handle_open(mock_ws)
 
         # Connection should NOT be marked as connected
         assert not ws_client._connected.is_set()
         # ws.close() should have been called
         mock_ws.close.assert_called_once()
-        # Error callback should have been invoked
-        error_cb.assert_called_once()
 
 
 # ---------------------------------------------------------------------------
@@ -1064,6 +1083,12 @@ def test_no_callback_uses_queue(self, ws_client) -> None:
         assert not ws_client._queue.empty()
         assert ws_client._queue.get_nowait() == {"event": "data"}
 
+    def test_receive_raises_when_message_callback_registered(self, ws_client) -> None:
+        ws_client.on_message(Mock())
+        ws_client._connected.set()
+        with pytest.raises(RuntimeError, match="on_message callback"):
+            list(ws_client.receive())
+
 
 # ---------------------------------------------------------------------------
 # disconnect(wait=...)
@@ -1084,6 +1109,26 @@ def test_disconnect_wait_blocks_until_thread_finishes(
         assert client._finished.is_set()
         assert not client._thread.is_alive()
 
+    def test_disconnect_wait_from_callback_does_not_self_join(
+        self, mock_session
+    ) -> None:
+        """disconnect(wait=True) from inside a callback must not raise."""
+        client = _MistWebsocket(mock_session, channels=["/ch"])
+        error_from_cb = []
+
+        def on_close_cb(code, msg):
+            try:
+                client.disconnect(wait=True, timeout=1)
+            except Exception as exc:
+                error_from_cb.append(exc)
+
+        client.on_close(on_close_cb)
+        mock_ws = Mock()
+        client._ws = mock_ws
+        client._run_forever_safe()
+
+        assert not error_from_cb, f"Callback raised: {error_from_cb}"
+
 
 # ---------------------------------------------------------------------------
 # Cookie edge cases
@@ -1131,3 +1176,130 @@ def test_session_with_url_rejects_http(self, mock_session) -> None:
     def test_session_with_url_accepts_wss(self, mock_session) -> None:
         ws = SessionWithUrl(mock_session, url="wss://api-ws.mist.com/stream")
         assert ws._build_ws_url() == "wss://api-ws.mist.com/stream"
+
+
+# ---------------------------------------------------------------------------
+# Connect / disconnect / connect cycle
+# ---------------------------------------------------------------------------
+
+
+class TestConnectDisconnectCycle:
+    """Verify that connect → disconnect → connect works cleanly."""
+
+    @patch("mistapi.websockets.__ws_client.websocket.WebSocketApp")
+    def test_reconnect_after_disconnect(self, mock_ws_cls, mock_session) -> None:
+        mock_ws_cls.return_value = Mock()
+        client = _MistWebsocket(mock_session, channels=["/ch"])
+
+        # First cycle
+        client.connect(run_in_background=True)
+        client.disconnect(wait=True, timeout=5)
+        assert client._finished.is_set()
+
+        # Second cycle — should not raise
+        client.connect(run_in_background=True)
+        client.disconnect(wait=True, timeout=5)
+        assert client._finished.is_set()
+
+
+# ---------------------------------------------------------------------------
+# receive() exits when thread dies (no sentinel scenario)
+# ---------------------------------------------------------------------------
+
+
+class TestReceiveFinishedExit:
+    """Verify receive() exits when _finished is set even without a sentinel."""
+
+    def test_receive_exits_when_finished_set_without_sentinel(
+        self, ws_client
+    ) -> None:
+        """Simulates a BaseException scenario where sentinel is never queued."""
+        ws_client._connected.set()
+        # Simulate: thread died, _finished set, _connected still set, no sentinel
+        ws_client._finished.set()
+        # receive() should exit promptly
+        results = list(ws_client.receive())
+        assert results == []
+
+    def test_receive_exits_when_finished_set_with_connected_cleared(
+        self, ws_client
+    ) -> None:
+        ws_client._connected.set()
+        ws_client._connected.clear()
+        ws_client._finished.set()
+        results = list(ws_client.receive())
+        assert results == []
+
+    def test_receive_drains_queue_when_connection_closed_before_receive(
+        self, ws_client
+    ) -> None:
+        """If connection opens, messages arrive, and connection closes before
+        receive() is called, the queued messages should still be yielded."""
+        ws_client._finished.clear()  # simulate connect() was called
+        # Simulate: connection opened, messages arrived, connection closed
+        ws_client._queue.put({"event": "a"})
+        ws_client._queue.put({"event": "b"})
+        ws_client._queue.put(None)  # sentinel from _run_forever_safe
+        ws_client._finished.set()  # thread finished
+        # _connected was set then cleared — currently unset
+
+        results = list(ws_client.receive())
+        assert results == [{"event": "a"}, {"event": "b"}]
+
+
+# ---------------------------------------------------------------------------
+# Blocking connect guard
+# ---------------------------------------------------------------------------
+
+
+class TestBlockingConnectGuard:
+    """Verify that the guard prevents double-connect in blocking mode."""
+
+    def test_connect_blocking_sets_finished_cleared(self, mock_session) -> None:
+        """_finished is cleared inside connect(), preventing concurrent connect."""
+        client = _MistWebsocket(mock_session, channels=["/ch"])
+        assert client._finished.is_set()  # starts set = ready
+
+        mock_ws = Mock()
+        with patch(
+            "mistapi.websockets.__ws_client.websocket.WebSocketApp",
+            return_value=mock_ws,
+        ):
+            client.connect(run_in_background=False)
+
+        # After _run_forever_safe returns, _finished is set again
+        assert client._finished.is_set()
+
+    def test_double_connect_raises_while_blocking(self, mock_session) -> None:
+        """If a blocking connect is in progress, a concurrent connect raises."""
+        client = _MistWebsocket(mock_session, channels=["/ch"])
+        barrier = threading.Event()
+
+        def blocking_run_forever(**kwargs):
+            barrier.wait(timeout=5)  # block until test releases
+
+        mock_ws = Mock()
+        mock_ws.run_forever.side_effect = blocking_run_forever
+
+        with patch(
+            "mistapi.websockets.__ws_client.websocket.WebSocketApp",
+            return_value=mock_ws,
+        ):
+            t = threading.Thread(
+                target=client.connect, kwargs={"run_in_background": False}
+            )
+            t.start()
+
+            # Wait for _finished to be cleared (connect entered)
+            for _ in range(50):
+                if not client._finished.is_set():
+                    break
+                threading.Event().wait(timeout=0.05)
+
+            # Second connect should raise
+            with pytest.raises(RuntimeError, match="Already connected"):
+                client.connect(run_in_background=True)
+
+            barrier.set()  # release blocking thread
+            t.join(timeout=5)
+            assert not t.is_alive()