fix: add a decorator to mark traits as mqtt only

Author: allenporter

roborock/devices/device_manager.py

@@ -149,7 +149,7 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
         match device.pv:
             case DeviceVersion.V1:
                 channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
-                trait = v1.create(product, channel.rpc_channel)
+                trait = v1.create(product, channel.rpc_channel, channel.mqtt_rpc_channel)
             case DeviceVersion.A01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 trait = a01.create(product, channel)

roborock/devices/traits/v1/__init__.py

@@ -1,6 +1,7 @@
 """Create traits for V1 devices."""
 
 from dataclasses import dataclass, field, fields
+import logging
 
 from roborock.containers import HomeData, HomeDataProduct
 from roborock.devices.traits import Trait
@@ -13,6 +14,8 @@
 from .status import StatusTrait
 from .volume import SoundVolumeTrait
 
+_LOGGER = logging.getLogger(__name__)
+
 __all__ = [
     "create",
     "PropertiesApi",
@@ -40,10 +43,12 @@ class PropertiesApi(Trait):
 
     # In the future optional fields can be added below based on supported features
 
-    def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None:
+    def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> None:
         """Initialize the V1TraitProps with None values."""
         self.status = StatusTrait(product)
+        #self.status._rpc_channel = rpc_channel
         self.maps = MapsTrait(self.status)
+        #self.maps._rpc_channel = mqtt_rpc_channel
 
         # This is a hack to allow setting the rpc_channel on all traits. This is
         # used so we can preserve the dataclass behavior when the values in the
@@ -53,9 +58,12 @@ def __init__(self, product: HomeDataProduct, rpc_channel: V1RpcChannel) -> None:
             if (trait := getattr(self, item.name, None)) is None:
                 trait = item.type()
                 setattr(self, item.name, trait)
-            trait._rpc_channel = rpc_channel
+            if hasattr(trait, "mqtt_rpc_channel"):  # @common.mqtt_rpc_channel
+                trait._rpc_channel = mqtt_rpc_channel
+            else:
+                trait._rpc_channel = rpc_channel
 
 
-def create(product: HomeDataProduct, rpc_channel: V1RpcChannel) -> PropertiesApi:
+def create(product: HomeDataProduct, rpc_channel: V1RpcChannel, mqtt_rpc_channel: V1RpcChannel) -> PropertiesApi:
     """Create traits for V1 devices."""
-    return PropertiesApi(product, rpc_channel)
+    return PropertiesApi(product, rpc_channel, mqtt_rpc_channel)

roborock/devices/traits/v1/common.py

@@ -5,12 +5,16 @@
 
 from abc import ABC
 from dataclasses import dataclass, fields
+from functools import wraps
+import logging
 from typing import ClassVar, Self
 
 from roborock.containers import RoborockBase
 from roborock.devices.v1_rpc_channel import V1RpcChannel
 from roborock.roborock_typing import RoborockCommand
 
+_LOGGER = logging.getLogger(__name__)
+
 V1ResponseData = dict | list | int | str
 
 
@@ -115,3 +119,16 @@ def _parse_response(cls, response: V1ResponseData) -> Self:
             raise ValueError(f"Unexpected response format: {response!r}")
         value_field = _get_value_field(cls)
         return cls(**{value_field: response})
+
+
+def mqtt_rpc_channel(cls):
+    """Decorator to mark a function as cloud only.
+    
+    Normally a trait uses an adaptive rpc channel that can use either local
+    or cloud communication depending on what is available. This will force
+    the trait to always use the cloud rpc channel.
+    """
+    def wrapper(*args, **kwargs):
+        return cls(*args, **kwargs)
+    cls.mqtt_rpc_channel = True  # type: ignore[attr-defined]
+    return wrapper

roborock/devices/traits/v1/maps.py

@@ -16,6 +16,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+@common.mqtt_rpc_channel
 class MapsTrait(MultiMapsList, common.V1TraitMixin):
     """Trait for managing the maps of Roborock devices.
 

tests/devices/test_v1_device.py

@@ -35,13 +35,18 @@ def rpc_channel_fixture() -> AsyncMock:
     return AsyncMock()
 
 
+@pytest.fixture(autouse=True, name="mqtt_rpc_channel")
+def mqtt_rpc_channel_fixture() -> AsyncMock:
+    """Fixture to set up the channel for tests."""
+    return AsyncMock()
+
 @pytest.fixture(autouse=True, name="device")
-def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock) -> RoborockDevice:
+def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel: AsyncMock) -> RoborockDevice:
     """Fixture to set up the device for tests."""
     return RoborockDevice(
         device_info=HOME_DATA.devices[0],
         channel=channel,
-        trait=v1.create(HOME_DATA.products[0], rpc_channel),
+        trait=v1.create(HOME_DATA.products[0], rpc_channel, mqtt_rpc_channel),
     )
 
 

tests/devices/traits/v1/fixtures.py

@@ -27,11 +27,17 @@ def rpc_channel_fixture() -> AsyncMock:
     return AsyncMock()
 
 
+
+@pytest.fixture(autouse=True, name="mock_mqtt_rpc_channel")
+def mqtt_rpc_channel_fixture() -> AsyncMock:
+    """Fixture to set up the channel for tests."""
+    return AsyncMock()
+
 @pytest.fixture(autouse=True, name="device")
-def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock) -> RoborockDevice:
+def device_fixture(channel: AsyncMock, mock_rpc_channel: AsyncMock, mock_mqtt_rpc_channel: AsyncMock) -> RoborockDevice:
     """Fixture to set up the device for tests."""
     return RoborockDevice(
         device_info=HOME_DATA.devices[0],
         channel=channel,
-        trait=v1.create(HOME_DATA.products[0], mock_rpc_channel),
+        trait=v1.create(HOME_DATA.products[0], mock_rpc_channel, mock_mqtt_rpc_channel),
     )

tests/devices/traits/v1/test_maps.py

@@ -58,15 +58,20 @@ def maps_trait(device: RoborockDevice) -> MapsTrait:
 async def test_refresh_maps_trait(
     maps_trait: MapsTrait,
     mock_rpc_channel: AsyncMock,
+    mock_mqtt_rpc_channel: AsyncMock,
     status_trait: StatusTrait,
 ) -> None:
     """Test successfully getting multi maps list."""
     # Setup mock to return the sample multi maps list
     mock_rpc_channel.send_command.side_effect = [
         mock_data.STATUS,  # Initial status fetch
+    ]
+    mock_mqtt_rpc_channel.send_command.side_effect = [
         MULTI_MAP_LIST_DATA,
     ]
     await status_trait.refresh()
+    assert status_trait.current_map == 0
+
     # Populating the status information gives us the current map
     # flag, but we have not loaded the rest of the information.
     assert maps_trait.current_map == 0
@@ -96,22 +101,27 @@ async def test_refresh_maps_trait(
     assert maps_trait.current_map_info.name == "Map 1"
 
     # Verify the RPC call was made correctly
-    assert mock_rpc_channel.send_command.call_count == 2
+    assert mock_rpc_channel.send_command.call_count == 1
     mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_STATUS)
-    mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST)
+    assert mock_mqtt_rpc_channel.send_command.call_count == 1
+    mock_mqtt_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST)
 
 
 async def test_set_current_map(
     status_trait: StatusTrait,
     maps_trait: MapsTrait,
     mock_rpc_channel: AsyncMock,
+    mock_mqtt_rpc_channel: AsyncMock,
 ) -> None:
     """Test successfully setting the current map."""
+    assert hasattr(maps_trait, "mqtt_rpc_channel")
     mock_rpc_channel.send_command.side_effect = [
         mock_data.STATUS,  # Initial status fetch
+        UPDATED_STATUS,  # Response for refreshing status
+    ]
+    mock_mqtt_rpc_channel.send_command.side_effect = [
         MULTI_MAP_LIST_DATA,  # Response for LOAD_MULTI_MAP
         {},  # Response for setting the current map
-        UPDATED_STATUS,  # Response for refreshing status
     ]
     await status_trait.refresh()
 
@@ -139,7 +149,8 @@ async def test_set_current_map(
     # 2. GET_MULTI_MAPS_LIST to get the map list
     # 3. LOAD_MULTI_MAP to set the map
     # 4. GET_STATUS to refresh the current map in status
-    assert mock_rpc_channel.send_command.call_count == 4
+    assert mock_rpc_channel.send_command.call_count == 2
     mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_STATUS)
-    mock_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST)
-    mock_rpc_channel.send_command.assert_any_call(RoborockCommand.LOAD_MULTI_MAP, params=[123])
+    assert mock_mqtt_rpc_channel.send_command.call_count == 2
+    mock_mqtt_rpc_channel.send_command.assert_any_call(RoborockCommand.GET_MULTI_MAPS_LIST)
+    mock_mqtt_rpc_channel.send_command.assert_any_call(RoborockCommand.LOAD_MULTI_MAP, params=[123])