chore: Hook up the trait to the device and CLI

Author: allenporter

roborock/cli.py

@@ -42,7 +42,7 @@
 from pyshark.packet.packet import Packet  # type: ignore
 
 from roborock import SHORT_MODEL_TO_ENUM, DeviceFeatures, RoborockCommand, RoborockException
-from roborock.containers import DeviceData, HomeData, NetworkInfo, RoborockBase, UserData
+from roborock.containers import CombinedMapInfo, DeviceData, HomeData, NetworkInfo, RoborockBase, UserData
 from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device import RoborockDevice
 from roborock.devices.device_manager import DeviceManager, create_device_manager, create_home_data_api
@@ -116,6 +116,7 @@ class ConnectionCache(RoborockBase):
     email: str
     home_data: HomeData | None = None
     network_info: dict[str, NetworkInfo] | None = None
+    home_cache: dict[int, CombinedMapInfo] | None = None
 
 
 class DeviceConnectionManager:
@@ -258,14 +259,21 @@ def finish_session(self) -> None:
 
     async def get(self) -> CacheData:
         """Get cached value."""
+        _LOGGER.debug("Getting cache data")
         connection_cache = self.cache_data()
-        return CacheData(home_data=connection_cache.home_data, network_info=connection_cache.network_info or {})
+        return CacheData(
+            home_data=connection_cache.home_data,
+            network_info=connection_cache.network_info or {},
+            home_cache=connection_cache.home_cache,
+        )
 
     async def set(self, value: CacheData) -> None:
         """Set value in the cache."""
+        _LOGGER.debug("Setting cache data")
         connection_cache = self.cache_data()
         connection_cache.home_data = value.home_data
         connection_cache.network_info = value.network_info
+        connection_cache.home_cache = value.home_cache
         self.update(connection_cache)
 
 
@@ -533,6 +541,42 @@ async def rooms(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.rooms)
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.option("--refresh", is_flag=True, default=False, help="Refresh status before discovery.")
+@click.pass_context
+@async_command
+async def home(ctx, device_id: str, refresh: bool):
+    """Discover and cache home layout (maps and rooms)."""
+    context: RoborockContext = ctx.obj
+    device_manager = await context.get_device_manager()
+    device = await device_manager.get_device(device_id)
+    if device.v1_properties is None:
+        raise RoborockException(f"Device {device.name} does not support V1 protocol")
+
+    # Ensure we have the latest status before discovery
+    await device.v1_properties.status.refresh()
+
+    home_trait = device.v1_properties.home
+    await home_trait.discover_home()
+    if refresh:
+        await home_trait.refresh()
+
+    # Display the discovered home cache
+    if home_trait.home_cache:
+        cache_summary = {
+            map_flag: {
+                "name": map_data.name,
+                "room_count": len(map_data.rooms),
+                "rooms": [{"segment_id": room.segment_id, "name": room.name} for room in map_data.rooms],
+            }
+            for map_flag, map_data in home_trait.home_cache.items()
+        }
+        click.echo(dump_json(cache_summary))
+    else:
+        click.echo("No maps discovered")
+
+
 @click.command()
 @click.option("--device_id", required=True)
 @click.option("--cmd", required=True)
@@ -780,6 +824,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(consumables)
 cli.add_command(reset_consumable)
 cli.add_command(rooms)
+cli.add_command(home)
 
 
 def main():

roborock/devices/device_manager.py

@@ -157,6 +157,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     channel.rpc_channel,
                     channel.mqtt_rpc_channel,
                     channel.map_rpc_channel,
+                    cache,
                     map_parser_config=map_parser_config,
                 )
             case DeviceVersion.A01:

roborock/devices/traits/v1/__init__.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass, field, fields
 
 from roborock.containers import HomeData, HomeDataProduct
+from roborock.devices.cache import Cache
 from roborock.devices.traits import Trait
 from roborock.devices.v1_rpc_channel import V1RpcChannel
 from roborock.map.map_parser import MapParserConfig
@@ -12,6 +13,7 @@
 from .common import V1TraitMixin
 from .consumeable import ConsumableTrait
 from .do_not_disturb import DoNotDisturbTrait
+from .home import HomeTrait
 from .map_content import MapContentTrait
 from .maps import MapsTrait
 from .rooms import RoomsTrait
@@ -30,6 +32,7 @@
     "MapsTrait",
     "MapContentTrait",
     "ConsumableTrait",
+    "HomeTrait",
 ]
 
 
@@ -49,6 +52,7 @@ class PropertiesApi(Trait):
     maps: MapsTrait
     map_content: MapContentTrait
     consumables: ConsumableTrait
+    home: HomeTrait
 
     # In the future optional fields can be added below based on supported features
 
@@ -59,13 +63,15 @@ def __init__(
         rpc_channel: V1RpcChannel,
         mqtt_rpc_channel: V1RpcChannel,
         map_rpc_channel: V1RpcChannel,
+        cache: Cache,
         map_parser_config: MapParserConfig | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
         self.status = StatusTrait(product)
         self.rooms = RoomsTrait(home_data)
         self.maps = MapsTrait(self.status)
         self.map_content = MapContentTrait(map_parser_config)
+        self.home = HomeTrait(self.maps, self.rooms, cache)
         # 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
         # traits are updated, but still want to allow them to have a reference
@@ -90,7 +96,8 @@ def create(
     rpc_channel: V1RpcChannel,
     mqtt_rpc_channel: V1RpcChannel,
     map_rpc_channel: V1RpcChannel,
+    cache: Cache,
     map_parser_config: MapParserConfig | None = None,
 ) -> PropertiesApi:
     """Create traits for V1 devices."""
-    return PropertiesApi(product, home_data, rpc_channel, mqtt_rpc_channel, map_rpc_channel, map_parser_config)
+    return PropertiesApi(product, home_data, rpc_channel, mqtt_rpc_channel, map_rpc_channel, cache, map_parser_config)

roborock/devices/traits/v1/home.py

@@ -70,6 +70,7 @@ async def discover_home(self) -> None:
             raise RoborockException("Cannot perform home discovery without current map info")
 
         home_cache = await self._build_home_cache()
+        _LOGGER.debug("Home discovery complete, caching data for %d maps", len(home_cache))
         cache_data = await self._cache.get()
         cache_data.home_cache = home_cache
         await self._cache.set(cache_data)

tests/devices/test_v1_device.py

@@ -8,6 +8,7 @@
 from syrupy import SnapshotAssertion
 
 from roborock.containers import HomeData, S7MaxVStatus, UserData
+from roborock.devices.cache import NoCache
 from roborock.devices.device import RoborockDevice
 from roborock.devices.traits import v1
 from roborock.devices.traits.v1.common import V1TraitMixin
@@ -53,7 +54,7 @@ def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel:
     return RoborockDevice(
         device_info=HOME_DATA.devices[0],
         channel=channel,
-        trait=v1.create(HOME_DATA.products[0], HOME_DATA, rpc_channel, mqtt_rpc_channel, map_rpc_channel=AsyncMock()),
+        trait=v1.create(HOME_DATA.products[0], HOME_DATA, rpc_channel, mqtt_rpc_channel, AsyncMock(), NoCache()),
     )
 
 

tests/devices/traits/v1/fixtures.py

@@ -5,6 +5,7 @@
 import pytest
 
 from roborock.containers import HomeData, S7MaxVStatus, UserData
+from roborock.devices.cache import NoCache
 from roborock.devices.device import RoborockDevice
 from roborock.devices.traits import v1
 
@@ -53,5 +54,6 @@ def device_fixture(
             mock_rpc_channel,
             mock_mqtt_rpc_channel,
             mock_map_rpc_channel,
+            NoCache(),
         ),
     )