fix: Update mqtt channel to correctly handle multiple subscribers

Author: allenporter

roborock/devices/mqtt_channel.py

@@ -30,12 +30,14 @@ def __init__(self, mqtt_session: MqttSession, duid: str, local_key: str, rriot:
 
         self._decoder = create_mqtt_decoder(local_key)
         self._encoder = create_mqtt_encoder(local_key)
-        self._mqtt_unsub: Callable[[], None] | None = None
 
     @property
     def is_connected(self) -> bool:
-        """Return true if the channel is connected."""
-        return (self._mqtt_unsub is not None) and self._mqtt_session.connected
+        """Return true if the channel is connected.
+
+        This passes through the underlying MQTT session's connected state.
+        """
+        return self._mqtt_session.connected
 
     @property
     def _publish_topic(self) -> str:
@@ -52,9 +54,6 @@ async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callab
 
         The callback will be called with the message payload when a message is received.
 
-        All messages received will be processed through the provided callback, even
-        those sent in response to the `send_command` command.
-
         Returns a callable that can be used to unsubscribe from the topic.
         """
 
@@ -69,14 +68,7 @@ def message_handler(payload: bytes) -> None:
                 except Exception as e:
                     _LOGGER.exception("Uncaught error in message handler callback: %s", e)
 
-        self._mqtt_unsub = await self._mqtt_session.subscribe(self._subscribe_topic, message_handler)
-
-        def unsub_wrapper() -> None:
-            if self._mqtt_unsub is not None:
-                self._mqtt_unsub()
-                self._mqtt_unsub = None
-
-        return unsub_wrapper
+        return await self._mqtt_session.subscribe(self._subscribe_topic, message_handler)
 
     async def publish(self, message: RoborockMessage) -> None:
         """Publish a command message.

roborock/devices/v1_channel.py

@@ -79,7 +79,7 @@ def is_local_connected(self) -> bool:
     @property
     def is_mqtt_connected(self) -> bool:
         """Return whether MQTT connection is available."""
-        return self._mqtt_unsub is not None
+        return self._mqtt_unsub is not None and self._mqtt_channel.is_connected
 
     @property
     def rpc_channel(self) -> V1RpcChannel:

tests/devices/test_mqtt_channel.py

@@ -3,7 +3,7 @@
 import asyncio
 import json
 import logging
-from collections.abc import Callable, Generator
+from collections.abc import AsyncGenerator, Callable, Generator
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
@@ -63,23 +63,29 @@ def setup_mqtt_channel(mqtt_session: Mock) -> MqttChannel:
     )
 
 
-@pytest.fixture(name="received_messages", autouse=True)
-async def setup_subscribe_callback(mqtt_channel: MqttChannel) -> list[RoborockMessage]:
+@pytest.fixture(name="mqtt_subscribers", autouse=True)
+async def setup_subscribe_callback(mqtt_session: Mock) -> AsyncGenerator[list[Callable[[bytes], None]], None]:
     """Fixture to record messages received by the subscriber."""
-    messages: list[RoborockMessage] = []
-    await mqtt_channel.subscribe(messages.append)
-    return messages
+    subscriber_callbacks = []
+
+    def mock_subscribe(_: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
+        subscriber_callbacks.append(callback)
+        return lambda: subscriber_callbacks.remove(callback)
+
+    mqtt_session.subscribe.side_effect = mock_subscribe
+    yield subscriber_callbacks
+    assert not subscriber_callbacks, "Not all subscribers were unsubscribed"
 
 
 @pytest.fixture(name="mqtt_message_handler")
-async def setup_message_handler(mqtt_session: Mock, mqtt_channel: MqttChannel) -> Callable[[bytes], None]:
+async def setup_message_handler(mqtt_subscribers: list[Callable[[bytes], None]]) -> Callable[[bytes], None]:
     """Fixture to allow simulating incoming MQTT messages."""
-    # Subscribe to set up message handling. We grab the message handler callback
-    # and use it to simulate receiving a response.
-    assert mqtt_session.subscribe
-    subscribe_call_args = mqtt_session.subscribe.call_args
-    message_handler = subscribe_call_args[0][1]
-    return message_handler
+
+    def invoke_all_callbacks(message: bytes) -> None:
+        for callback in mqtt_subscribers:
+            callback(message)
+
+    return invoke_all_callbacks
 
 
 @pytest.fixture
@@ -106,23 +112,6 @@ async def mock_home_data() -> HomeData:
     return HomeData.from_dict(mock_data.HOME_DATA_RAW)
 
 
-async def test_mqtt_channel(mqtt_session: Mock, mqtt_channel: MqttChannel) -> None:
-    """Test MQTT channel setup."""
-
-    unsub = Mock()
-    mqtt_session.subscribe.return_value = unsub
-
-    callback = Mock()
-    result = await mqtt_channel.subscribe(callback)
-
-    assert mqtt_session.subscribe.called
-    assert mqtt_session.subscribe.call_args[0][0] == "rr/m/o/user123/username/abc123"
-
-    unsub.assert_not_called()
-    result()
-    unsub.assert_called_once()
-
-
 async def test_publish_success(
     mqtt_session: Mock,
     mqtt_channel: MqttChannel,
@@ -148,13 +137,126 @@ async def test_publish_success(
 
 
 async def test_message_decode_error(
+    mqtt_channel: MqttChannel,
     mqtt_message_handler: Callable[[bytes], None],
     caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test an error during message decoding."""
+    callback = Mock()
+    unsub = await mqtt_channel.subscribe(callback)
+
     mqtt_message_handler(b"invalid_payload")
     await asyncio.sleep(0.01)  # yield
 
     assert len(caplog.records) == 1
     assert caplog.records[0].levelname == "WARNING"
     assert "Failed to decode MQTT message" in caplog.records[0].message
+    unsub()
+
+
+async def test_concurrent_subscribers(mqtt_session: Mock, mqtt_channel: MqttChannel) -> None:
+    """Test multiple concurrent subscribers receive all messages."""
+    # Set up multiple subscribers
+    subscriber1_messages: list[RoborockMessage] = []
+    subscriber2_messages: list[RoborockMessage] = []
+    subscriber3_messages: list[RoborockMessage] = []
+
+    unsub1 = await mqtt_channel.subscribe(subscriber1_messages.append)
+    unsub2 = await mqtt_channel.subscribe(subscriber2_messages.append)
+    unsub3 = await mqtt_channel.subscribe(subscriber3_messages.append)
+
+    # Verify that each subscription creates a separate call to the MQTT session
+    assert mqtt_session.subscribe.call_count == 3
+
+    # All subscriptions should be to the same topic
+    for call in mqtt_session.subscribe.call_args_list:
+        assert call[0][0] == "rr/m/o/user123/username/abc123"
+
+    # Get the message handlers for each subscriber
+    handler1 = mqtt_session.subscribe.call_args_list[0][0][1]
+    handler2 = mqtt_session.subscribe.call_args_list[1][0][1]
+    handler3 = mqtt_session.subscribe.call_args_list[2][0][1]
+
+    # Simulate receiving messages - each handler should decode the message independently
+    handler1(ENCODER(TEST_REQUEST))
+    handler2(ENCODER(TEST_REQUEST))
+    handler3(ENCODER(TEST_REQUEST))
+    await asyncio.sleep(0.01)  # yield
+
+    # All subscribers should receive the message
+    assert len(subscriber1_messages) == 1
+    assert len(subscriber2_messages) == 1
+    assert len(subscriber3_messages) == 1
+    assert subscriber1_messages[0] == TEST_REQUEST
+    assert subscriber2_messages[0] == TEST_REQUEST
+    assert subscriber3_messages[0] == TEST_REQUEST
+
+    # Send another message to all handlers
+    handler1(ENCODER(TEST_RESPONSE))
+    handler2(ENCODER(TEST_RESPONSE))
+    handler3(ENCODER(TEST_RESPONSE))
+    await asyncio.sleep(0.01)  # yield
+
+    # All subscribers should have received both messages
+    assert len(subscriber1_messages) == 2
+    assert len(subscriber2_messages) == 2
+    assert len(subscriber3_messages) == 2
+    assert subscriber1_messages == [TEST_REQUEST, TEST_RESPONSE]
+    assert subscriber2_messages == [TEST_REQUEST, TEST_RESPONSE]
+    assert subscriber3_messages == [TEST_REQUEST, TEST_RESPONSE]
+
+    # Test unsubscribing one subscriber
+    unsub1()
+
+    # Send another message only to remaining handlers
+    handler2(ENCODER(TEST_REQUEST2))
+    handler3(ENCODER(TEST_REQUEST2))
+    await asyncio.sleep(0.01)  # yield
+
+    # First subscriber should not have received the new message
+    assert len(subscriber1_messages) == 2
+    assert len(subscriber2_messages) == 3
+    assert len(subscriber3_messages) == 3
+    assert subscriber2_messages[2] == TEST_REQUEST2
+    assert subscriber3_messages[2] == TEST_REQUEST2
+
+    # Unsubscribe remaining subscribers
+    unsub2()
+    unsub3()
+
+
+async def test_concurrent_subscribers_with_callback_exception(
+    mqtt_session: Mock, mqtt_channel: MqttChannel, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test that exception in one subscriber callback doesn't affect others."""
+    caplog.set_level(logging.ERROR)
+
+    def failing_callback(message: RoborockMessage) -> None:
+        raise ValueError("Callback error")
+
+    subscriber2_messages: list[RoborockMessage] = []
+
+    unsub1 = await mqtt_channel.subscribe(failing_callback)
+    unsub2 = await mqtt_channel.subscribe(subscriber2_messages.append)
+
+    # Get the message handlers
+    handler1 = mqtt_session.subscribe.call_args_list[0][0][1]
+    handler2 = mqtt_session.subscribe.call_args_list[1][0][1]
+
+    # Simulate receiving a message - first handler will raise exception
+    handler1(ENCODER(TEST_REQUEST))
+    handler2(ENCODER(TEST_REQUEST))
+    await asyncio.sleep(0.01)  # yield
+
+    # Exception should be logged but other subscribers should still work
+    assert len(subscriber2_messages) == 1
+    assert subscriber2_messages[0] == TEST_REQUEST
+
+    # Check that exception was logged
+    error_records = [record for record in caplog.records if record.levelname == "ERROR"]
+    assert len(error_records) == 1
+    assert "Uncaught error in message handler callback" in error_records[0].message
+
+    # Unsubscribe all remaining subscribers
+    unsub1()
+    unsub2()

tests/devices/test_v1_channel.py

@@ -63,13 +63,15 @@
 
 
 @pytest.fixture(name="mock_mqtt_channel")
-def setup_mock_mqtt_channel() -> FakeChannel:
+async def setup_mock_mqtt_channel() -> FakeChannel:
     """Mock MQTT channel for testing."""
-    return FakeChannel()
+    channel = FakeChannel()
+    await channel.connect()
+    return channel
 
 
 @pytest.fixture(name="mock_local_channel")
-def setup_mock_local_channel() -> FakeChannel:
+async def setup_mock_local_channel() -> FakeChannel:
     """Mock Local channel for testing."""
     return FakeChannel()
 
@@ -150,6 +152,39 @@ async def test_v1_channel_subscribe_mqtt_only_success(
     assert not mock_mqtt_channel.subscribers
 
 
+async def test_v1_channel_mqtt_disconnected(
+    v1_channel: V1Channel,
+    mock_mqtt_channel: FakeChannel,
+    mock_local_session: Mock,
+    mock_local_channel: FakeChannel,
+) -> None:
+    """Test successful subscription with MQTT only (local connection fails)."""
+    # Setup: MQTT succeeds, local fails
+    mock_mqtt_channel.response_queue.append(TEST_NETWORK_INFO_RESPONSE)
+    mock_local_channel.connect.side_effect = RoborockException("Connection failed")
+
+    callback = Mock()
+    unsub = await v1_channel.subscribe(callback)
+
+    # Verify MQTT connection was established
+    assert mock_mqtt_channel.subscribers
+
+    # Verify local connection was attempted but failed
+    mock_local_session.assert_called_once_with(TEST_HOST)
+    mock_local_channel.connect.assert_called_once()
+
+    # Simulate an MQTT disconnection where the channel is not healthy
+    await mock_mqtt_channel.close()
+
+    # Verify properties
+    assert not v1_channel.is_mqtt_connected
+    assert not v1_channel.is_local_connected
+
+    # Test unsubscribe
+    unsub()
+    assert not mock_mqtt_channel.subscribers
+
+
 async def test_v1_channel_subscribe_both_connections_success(
     v1_channel: V1Channel,
     mock_mqtt_channel: Mock,