Clean up and accurately report WebSocket receive-loop termination (#102)

A recv() failure after the initial handshake now emits one
receiveMessageError and completes the connection wrapper without
re-raising through setup error handling, so it is no longer
relabeled as createConnectionError. Every receive-loop termination
(receive failure, cancellation, or intentional closure) marks the
client not ready and cancels its heartbeat task; intentional
closure emits no receive error. createConnectionError remains
reserved for connection establishment and the initial handshake.

Author: tylerlong

ringcentral/websocket/web_socket_client.py

@@ -15,6 +15,7 @@ def __init__(self, platform):
         self._web_socket = None
         self._done = False
         self._is_ready = False
+        self._heartbeat_task = None
         self._send_attempt_counter = 0
 
     async def create_new_connection(self):
@@ -81,38 +82,50 @@ async def open_connection(self, ws_uri, ws_access_token):
                 - Triggers the connectionCreated event upon successful connection establishment.
                 - Listens for incoming messages and triggers the receiveMessage event for each received message.
                 - Isolates receive-message handler failures so a handler that raises does not stop delivery to the remaining handlers or reception of future messages; each failure triggers the receiveMessageError event with the original exception.
-                - Triggers the createConnectionError event if an error occurs during the connection process and raises the exception.
+                - Triggers the createConnectionError event if an error occurs while establishing the connection or during the initial handshake, and raises the exception.
+                - A failure from receiving messages after the initial handshake triggers the receiveMessageError event once; it is not reported as a connection-creation error and does not propagate to the connection-setup or recovery wrappers.
+                - When the receive loop terminates (receive failure, cancellation, or intentional closure), the client is marked not ready and its heartbeat task is cancelled.
+                - An intentional closure does not trigger the receiveMessageError event.
         """
         try:
             websocket = await websockets.connect(
                 f"{ws_uri}?access_token={ws_access_token}"
             )
             connectionMessage = await websocket.recv()
-            connection_info = {}
-            connection_info["connection"] = websocket
-            connection_info["connection_details"] = connectionMessage
-            self._web_socket = connection_info
-            self._is_ready = True
-            self.trigger(WebSocketEvents.connectionCreated, self)
-
-            # heartbeat every 10 minutes
-            async def timer_function():
-                while True:
-                    if self._done:
-                        timer.cancel()
-                        break
-                    await asyncio.sleep(600)
-                    await self.send_message([{"type": "Heartbeat", "messageId": str(uuid.uuid4())}])
-            timer = asyncio.create_task(timer_function())
+        except Exception as e:
+            self.trigger(WebSocketEvents.createConnectionError, e)
+            raise
+
+        connection_info = {}
+        connection_info["connection"] = websocket
+        connection_info["connection_details"] = connectionMessage
+        self._web_socket = connection_info
+        self._is_ready = True
+        self.trigger(WebSocketEvents.connectionCreated, self)
+
+        # heartbeat every 10 minutes
+        async def timer_function():
+            while True:
+                if self._done:
+                    timer.cancel()
+                    break
+                await asyncio.sleep(600)
+                await self.send_message([{"type": "Heartbeat", "messageId": str(uuid.uuid4())}])
+        timer = asyncio.create_task(timer_function())
+        self._heartbeat_task = timer
 
+        try:
             await asyncio.sleep(0)
             while True:
                 message = await websocket.recv()
                 self.trigger(WebSocketEvents.receiveMessage, message)
                 await asyncio.sleep(0)
         except Exception as e:
-            self.trigger(WebSocketEvents.createConnectionError, e)
-            raise
+            if not self._done:
+                self.trigger(WebSocketEvents.receiveMessageError, e)
+        finally:
+            self._is_ready = False
+            timer.cancel()
 
     def trigger(self, event, *args, **kw):
         """

ringcentral/websocket/web_socket_client_test.py

@@ -30,16 +30,27 @@ def __init__(self):
         self.sent = []
         self.received = []
         self.closed = False
+        self._fail_next_recv = False
+        self._receive_error = None
         self._frames = asyncio.Queue()
 
+    def fail_next_receive(self, error):
+        self._fail_next_recv = True
+        self._receive_error = error
+
     def push(self, frame):
         self._frames.put_nowait(frame)
 
     def has_pending_frames(self):
         return self._frames.qsize() > 0
 
     async def recv(self):
+        if self._fail_next_recv:
+            self._fail_next_recv = False
+            raise self._receive_error
         frame = await self._frames.get()
+        if frame is None:
+            raise ConnectionResetError("connection closed")
         self.received.append(frame)
         return frame
 
@@ -48,6 +59,7 @@ async def send(self, message):
 
     async def close(self):
         self.closed = True
+        self._frames.put_nowait(None)
 
 
 class RecordingHandler:
@@ -294,6 +306,121 @@ async def test_unrelated_event_dispatch_semantics_remain_unchanged(self):
         self.assertEqual(failing.calls, [(self.client,)])
         self.assertEqual(healthy.calls, [])
 
+    async def test_recv_failure_after_handshake_reports_receive_error_and_cleans_up(self):
+        receive_error = RuntimeError("receive failure")
+        receive_errors = RecordingHandler()
+        connection_errors = RecordingHandler()
+        self.client.on(WebSocketEvents.receiveMessageError, receive_errors)
+        self.client.on(WebSocketEvents.createConnectionError, connection_errors)
+
+        sockets = []
+
+        def on_connected(*_):
+            sockets[0].fail_next_receive(receive_error)
+
+        self.client.on(WebSocketEvents.connectionCreated, on_connected)
+
+        fake_socket, connection_task = await self.start_receive_loop()
+        fake_socket.push("[/heartbeat]")
+        sockets.append(fake_socket)
+
+        result = await connection_task
+
+        self.assertIsNone(result)
+        self.assertEqual(len(receive_errors.calls), 1)
+        self.assertEqual(len(receive_errors.calls[0]), 1)
+        self.assertIs(receive_errors.calls[0][0], receive_error)
+        self.assertEqual(connection_errors.calls, [])
+        self.assertFalse(self.client._is_ready)
+        await self.wait_until(lambda: self.client._heartbeat_task.done())
+        self.assertTrue(self.client._heartbeat_task.cancelled())
+
+    async def test_receive_loop_cancellation_marks_client_not_ready_and_cancels_heartbeat(self):
+        receive_errors = RecordingHandler()
+        connection_errors = RecordingHandler()
+        self.client.on(WebSocketEvents.receiveMessageError, receive_errors)
+        self.client.on(WebSocketEvents.createConnectionError, connection_errors)
+
+        fake_socket, connection_task = await self.start_receive_loop()
+        fake_socket.push("[/heartbeat]")
+        await self.wait_until(lambda: self.client._is_ready)
+        self.assertTrue(self.client._is_ready)
+
+        await self.stop_receive_loop(connection_task)
+
+        self.assertFalse(self.client._is_ready)
+        await self.wait_until(lambda: self.client._heartbeat_task.done())
+        self.assertTrue(self.client._heartbeat_task.cancelled())
+        self.assertEqual(receive_errors.calls, [])
+        self.assertEqual(connection_errors.calls, [])
+
+    async def test_intentional_close_performs_cleanup_without_receive_error(self):
+        receive_errors = RecordingHandler()
+        connection_errors = RecordingHandler()
+        close_errors = RecordingHandler()
+        self.client.on(WebSocketEvents.receiveMessageError, receive_errors)
+        self.client.on(WebSocketEvents.createConnectionError, connection_errors)
+        self.client.on(WebSocketEvents.closeConnectionError, close_errors)
+
+        fake_socket, connection_task = await self.start_receive_loop()
+        fake_socket.push("[/heartbeat]")
+        await self.wait_until(lambda: self.client._is_ready)
+
+        await self.client.close_connection()
+        result = await connection_task
+
+        self.assertIsNone(result)
+        self.assertTrue(fake_socket.closed)
+        self.assertFalse(self.client._is_ready)
+        await self.wait_until(lambda: self.client._heartbeat_task.done())
+        self.assertTrue(self.client._heartbeat_task.cancelled())
+        self.assertEqual(receive_errors.calls, [])
+        self.assertEqual(connection_errors.calls, [])
+        self.assertEqual(close_errors.calls, [])
+
+    async def test_connection_establishment_failure_retains_create_connection_error(self):
+        connect_error = RuntimeError("connect failure")
+        receive_errors = RecordingHandler()
+        connection_errors = RecordingHandler()
+        self.client.on(WebSocketEvents.receiveMessageError, receive_errors)
+        self.client.on(WebSocketEvents.createConnectionError, connection_errors)
+
+        async def connect(uri):
+            raise connect_error
+
+        with mock.patch("websockets.connect", connect):
+            with self.assertRaises(RuntimeError):
+                await self.client.create_new_connection()
+
+        self.assertTrue(connection_errors.calls)
+        for call in connection_errors.calls:
+            self.assertIs(call[0], connect_error)
+        self.assertEqual(receive_errors.calls, [])
+
+    async def test_initial_handshake_failure_retains_create_connection_error(self):
+        handshake_error = RuntimeError("handshake failure")
+        receive_errors = RecordingHandler()
+        connection_errors = RecordingHandler()
+        self.client.on(WebSocketEvents.receiveMessageError, receive_errors)
+        self.client.on(WebSocketEvents.createConnectionError, connection_errors)
+
+        fake_socket = FakeWebSocket()
+        fake_socket.fail_next_receive(handshake_error)
+
+        async def connect(uri):
+            return fake_socket
+
+        with mock.patch("websockets.connect", connect):
+            with self.assertRaises(RuntimeError):
+                await self.client.create_new_connection()
+
+        self.assertTrue(connection_errors.calls)
+        for call in connection_errors.calls:
+            self.assertIs(call[0], handshake_error)
+        self.assertEqual(receive_errors.calls, [])
+        self.assertFalse(self.client._is_ready)
+        self.assertIsNone(self.client._heartbeat_task)
+
 
 if __name__ == "__main__":
     unittest.main()