fix: Fix bugs in the subscription idle timeout

We need to actually track which subscriptionsare active and don't add duplicate subscriptions in those cases. We keep a separate object to track the subscription state since it is different than the callback logic (e.g. subscribe callback is removed from the list when unsubscribe happens)

Author: allenporter

roborock/mqtt/roborock_session.py

@@ -69,6 +69,7 @@ def __init__(
         self._stop = False
         self._backoff = MIN_BACKOFF_INTERVAL
         self._client: aiomqtt.Client | None = None
+        self._client_subscribed_topics: set[str] = set()
         self._client_lock = asyncio.Lock()
         self._listeners: CallbackMap[str, bytes] = CallbackMap(_LOGGER)
         self._connection_task: asyncio.Task[None] | None = None
@@ -218,7 +219,7 @@ async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
                 # Re-establish any existing subscriptions
                 async with self._client_lock:
                     self._client = client
-                    for topic in self._listeners.keys():
+                    for topic in self._client_subscribed_topics:
                         _LOGGER.debug("Re-establishing subscription to topic %s", topic)
                         # TODO: If this fails it will break the whole connection. Make
                         # this retry again in the background with backoff.
@@ -249,32 +250,42 @@ async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Call
         unsub = self._listeners.add_callback(topic, callback)
 
         async with self._client_lock:
-            if self._client:
-                _LOGGER.debug("Establishing subscription to topic %s", topic)
-                try:
-                    await self._client.subscribe(topic)
-                except MqttError as err:
-                    # Clean up the callback if subscription fails
-                    unsub()
-                    raise MqttSessionException(f"Error subscribing to topic: {err}") from err
-            else:
-                _LOGGER.debug("Client not connected, will establish subscription later")
-
-        def schedule_unsubscribe():
+            if topic not in self._client_subscribed_topics:
+                self._client_subscribed_topics.add(topic)
+                if self._client:
+                    _LOGGER.debug("Establishing subscription to topic %s", topic)
+                    try:
+                        await self._client.subscribe(topic)
+                    except MqttError as err:
+                        # Clean up the callback if subscription fails
+                        unsub()
+                        self._client_subscribed_topics.discard(topic)
+                        raise MqttSessionException(f"Error subscribing to topic: {err}") from err
+                else:
+                    _LOGGER.debug("Client not connected, will establish subscription later")
+
+        def schedule_unsubscribe() -> None:
             async def idle_unsubscribe():
                 try:
                     await asyncio.sleep(self._topic_idle_timeout.total_seconds())
                     # Only unsubscribe if there are no callbacks left for this topic
                     if not self._listeners.get_callbacks(topic):
                         async with self._client_lock:
+                            # Check again if we have listeners, in case a subscribe happened
+                            # while we were waiting for the lock or after we popped the timer.
+                            if self._listeners.get_callbacks(topic):
+                                _LOGGER.debug("Skipping unsubscribe for %s, new listeners added", topic)
+                                return
+
+                            self._idle_timers.pop(topic, None)
+                            self._client_subscribed_topics.discard(topic)
+
                             if self._client:
                                 _LOGGER.debug("Idle timeout expired, unsubscribing from topic %s", topic)
                                 try:
                                     await self._client.unsubscribe(topic)
                                 except MqttError as err:
                                     _LOGGER.warning("Error unsubscribing from topic %s: %s", topic, err)
-                    # Clean up timer from dict
-                    self._idle_timers.pop(topic, None)
                 except asyncio.CancelledError:
                     _LOGGER.debug("Idle unsubscribe for topic %s cancelled", topic)
 
@@ -286,7 +297,10 @@ def delayed_unsub():
             unsub()  # Remove the callback from CallbackMap
             # If no more callbacks for this topic, start idle timer
             if not self._listeners.get_callbacks(topic):
+                _LOGGER.debug("Unsubscribing topic %s, starting idle timer", topic)
                 schedule_unsubscribe()
+            else:
+                _LOGGER.debug("Unsubscribing topic %s, still have active callbacks", topic)
 
         return delayed_unsub
 

tests/mqtt/test_roborock_session.py

@@ -368,6 +368,47 @@ async def test_idle_timeout_multiple_callbacks(mock_mqtt_client: AsyncMock) -> N
     await session.close()
 
 
+async def test_subscription_reuse(mock_mqtt_client: AsyncMock) -> None:
+    """Test that subscriptions are reused and not duplicated."""
+    session = RoborockMqttSession(FAKE_PARAMS)
+    await session.start()
+    assert session.connected
+
+    # 1. First subscription
+    cb1 = Mock()
+    unsub1 = await session.subscribe("topic1", cb1)
+
+    # Verify subscribe called
+    mock_mqtt_client.subscribe.assert_called_with("topic1")
+    mock_mqtt_client.subscribe.reset_mock()
+
+    # 2. Second subscription (same topic)
+    cb2 = Mock()
+    unsub2 = await session.subscribe("topic1", cb2)
+
+    # Verify subscribe NOT called
+    mock_mqtt_client.subscribe.assert_not_called()
+
+    # 3. Unsubscribe one
+    unsub1()
+    # Verify unsubscribe NOT called (still have cb2)
+    mock_mqtt_client.unsubscribe.assert_not_called()
+
+    # 4. Unsubscribe second (starts idle timer)
+    unsub2()
+    # Verify unsubscribe NOT called yet (idle)
+    mock_mqtt_client.unsubscribe.assert_not_called()
+
+    # 5. Resubscribe during idle
+    cb3 = Mock()
+    _ = await session.subscribe("topic1", cb3)
+
+    # Verify subscribe NOT called (reused)
+    mock_mqtt_client.subscribe.assert_not_called()
+
+    await session.close()
+
+
 @pytest.mark.parametrize(
     ("side_effect", "expected_exception", "match"),
     [