fix: fix a bug with local / mqtt fallback (#475)

Author: allenporter

roborock/devices/v1_channel.py

@@ -22,7 +22,7 @@
 from .channel import Channel
 from .local_channel import LocalChannel, LocalSession, create_local_session
 from .mqtt_channel import MqttChannel
-from .v1_rpc_channel import V1RpcChannel, create_combined_rpc_channel, create_mqtt_rpc_channel
+from .v1_rpc_channel import PickFirstAvailable, V1RpcChannel, create_local_rpc_channel, create_mqtt_rpc_channel
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -60,7 +60,11 @@ def __init__(
         self._mqtt_rpc_channel = create_mqtt_rpc_channel(mqtt_channel, security_data)
         self._local_session = local_session
         self._local_channel: LocalChannel | None = None
-        self._combined_rpc_channel: V1RpcChannel | None = None
+        self._local_rpc_channel: V1RpcChannel | None = None
+        # Prefer local, fallback to MQTT
+        self._combined_rpc_channel = PickFirstAvailable(
+            [lambda: self._local_rpc_channel, lambda: self._mqtt_rpc_channel]
+        )
         self._mqtt_unsub: Callable[[], None] | None = None
         self._local_unsub: Callable[[], None] | None = None
         self._callback: Callable[[RoborockMessage], None] | None = None
@@ -84,7 +88,7 @@ def is_mqtt_connected(self) -> bool:
     @property
     def rpc_channel(self) -> V1RpcChannel:
         """Return the combined RPC channel prefers local with a fallback to MQTT."""
-        return self._combined_rpc_channel or self._mqtt_rpc_channel
+        return self._combined_rpc_channel
 
     @property
     def mqtt_rpc_channel(self) -> V1RpcChannel:
@@ -160,7 +164,7 @@ async def _local_connect(self) -> Callable[[], None]:
         except RoborockException as e:
             self._local_channel = None
             raise RoborockException(f"Error connecting to local device {self._device_uid}: {e}") from e
-        self._combined_rpc_channel = create_combined_rpc_channel(self._local_channel, self._mqtt_rpc_channel)
+        self._local_rpc_channel = create_local_rpc_channel(self._local_channel)
         return await self._local_channel.subscribe(self._on_local_message)
 
     def _on_mqtt_message(self, message: RoborockMessage) -> None:

roborock/devices/v1_rpc_channel.py

@@ -88,16 +88,15 @@ async def _send_raw_command(
         raise NotImplementedError
 
 
-class CombinedV1RpcChannel(BaseV1RpcChannel):
-    """A V1 RPC channel that can use both local and MQTT channels, preferring local when available."""
+class PickFirstAvailable(BaseV1RpcChannel):
+    """A V1 RPC channel that tries multiple channels and picks the first that works."""
 
     def __init__(
-        self, local_channel: LocalChannel, local_rpc_channel: V1RpcChannel, mqtt_channel: V1RpcChannel
+        self,
+        channel_cbs: list[Callable[[], V1RpcChannel | None]],
     ) -> None:
-        """Initialize the combined channel with local and MQTT channels."""
-        self._local_channel = local_channel
-        self._local_rpc_channel = local_rpc_channel
-        self._mqtt_rpc_channel = mqtt_channel
+        """Initialize the pick-first-available channel."""
+        self._channel_cbs = channel_cbs
 
     async def _send_raw_command(
         self,
@@ -106,9 +105,10 @@ async def _send_raw_command(
         params: ParamsType = None,
     ) -> Any:
         """Send a command and return a parsed response RoborockBase type."""
-        if self._local_channel.is_connected:
-            return await self._local_rpc_channel.send_command(method, params=params)
-        return await self._mqtt_rpc_channel.send_command(method, params=params)
+        for channel_cb in self._channel_cbs:
+            if channel := channel_cb():
+                return await channel.send_command(method, params=params)
+        raise RoborockException("No available connection to send command")
 
 
 class PayloadEncodedV1RpcChannel(BaseV1RpcChannel):
@@ -170,11 +170,10 @@ def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityDa
     )
 
 
-def create_combined_rpc_channel(local_channel: LocalChannel, mqtt_rpc_channel: V1RpcChannel) -> V1RpcChannel:
-    """Create a V1 RPC channel that combines local and MQTT channels."""
-    local_rpc_channel = PayloadEncodedV1RpcChannel(
+def create_local_rpc_channel(local_channel: LocalChannel) -> V1RpcChannel:
+    """Create a V1 RPC channel using a local channel."""
+    return PayloadEncodedV1RpcChannel(
         "local",
         local_channel,
         lambda x: x.encode_message(RoborockMessageProtocol.GENERAL_REQUEST),
     )
-    return CombinedV1RpcChannel(local_channel, local_rpc_channel, mqtt_rpc_channel)

tests/devices/test_v1_channel.py

@@ -497,6 +497,9 @@ async def test_v1_channel_full_subscribe_and_command_flow(
         local_session=mock_local_session,
         cache=InMemoryCache(),
     )
+    # Get a handle to the V1RpcChannel. It may change which connection is
+    # active, but getting now to reproduce a bug where it doesn't change.
+    rpc_channel = v1_channel.rpc_channel
 
     # Mock network info for local connection
     callback = Mock()
@@ -512,7 +515,7 @@ async def test_v1_channel_full_subscribe_and_command_flow(
 
     # Send a command (should use local)
     mock_local_channel.response_queue.append(TEST_RESPONSE)
-    result = await v1_channel.rpc_channel.send_command(
+    result = await rpc_channel.send_command(
         RoborockCommand.GET_STATUS,
         response_type=S5MaxStatus,
     )