feat: Add a trait for reading NetworkInfo from the device (#558)

* feat: Add a trait for reading NetworkInfo from the device

This is mirroring what is done during device discovery, but also allows it to be exposed for the caller to read network information for display purposes. This will always prefer reading from the cache.

* chore: Update test assertion for network info

Author: allenporter

roborock/cli.py

@@ -731,6 +731,16 @@ async def home(ctx, device_id: str, refresh: bool):
         click.echo("No maps discovered")
 
 
+@session.command()
+@click.option("--device_id", required=True)
+@click.pass_context
+@async_command
+async def network_info(ctx, device_id: str):
+    """Get device network information."""
+    context: RoborockContext = ctx.obj
+    await _display_v1_trait(context, device_id, lambda v1: v1.network_info)
+
+
 @click.command()
 @click.option("--device_id", required=True)
 @click.option("--cmd", required=True)
@@ -979,6 +989,7 @@ def write_markdown_table(product_features: dict[str, dict[str, any]], all_featur
 cli.add_command(dnd)
 cli.add_command(flow_led_status)
 cli.add_command(led_status)
+cli.add_command(network_info)
 
 
 def main():

roborock/devices/device_manager.py

@@ -157,6 +157,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
             case DeviceVersion.V1:
                 channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
                 trait = v1.create(
+                    device.duid,
                     product,
                     home_data,
                     channel.rpc_channel,

roborock/devices/traits/v1/__init__.py

@@ -54,6 +54,7 @@
 from .led_status import LedStatusTrait
 from .map_content import MapContentTrait
 from .maps import MapsTrait
+from .network_info import NetworkInfoTrait
 from .rooms import RoomsTrait
 from .smart_wash_params import SmartWashParamsTrait
 from .status import StatusTrait
@@ -83,6 +84,7 @@
     "DustCollectionModeTrait",
     "WashTowelModeTrait",
     "SmartWashParamsTrait",
+    "NetworkInfoTrait",
 ]
 
 
@@ -105,6 +107,7 @@ class PropertiesApi(Trait):
     consumables: ConsumableTrait
     home: HomeTrait
     device_features: DeviceFeaturesTrait
+    network_info: NetworkInfoTrait
 
     # Optional features that may not be supported on all devices
     child_lock: ChildLockTrait | None = None
@@ -117,6 +120,7 @@ class PropertiesApi(Trait):
 
     def __init__(
         self,
+        device_uid: str,
         product: HomeDataProduct,
         home_data: HomeData,
         rpc_channel: V1RpcChannel,
@@ -126,6 +130,7 @@ def __init__(
         map_parser_config: MapParserConfig | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
+        self._device_uid = device_uid
         self._rpc_channel = rpc_channel
         self._mqtt_rpc_channel = mqtt_rpc_channel
         self._map_rpc_channel = map_rpc_channel
@@ -138,6 +143,7 @@ def __init__(
         self.map_content = MapContentTrait(map_parser_config)
         self.home = HomeTrait(self.status, self.maps, self.rooms, cache)
         self.device_features = DeviceFeaturesTrait(product.product_nickname, cache)
+        self.network_info = NetworkInfoTrait(device_uid, cache)
 
         # Dynamically create any traits that need to be populated
         for item in fields(self):
@@ -243,6 +249,7 @@ async def _set_cached_trait_data(self, name: str, value: Any) -> None:
 
 
 def create(
+    device_uid: str,
     product: HomeDataProduct,
     home_data: HomeData,
     rpc_channel: V1RpcChannel,
@@ -252,4 +259,13 @@ def create(
     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, cache, map_parser_config)
+    return PropertiesApi(
+        device_uid,
+        product,
+        home_data,
+        rpc_channel,
+        mqtt_rpc_channel,
+        map_rpc_channel,
+        cache,
+        map_parser_config,
+    )

roborock/devices/traits/v1/network_info.py

@@ -0,0 +1,56 @@
+"""Trait for device network information."""
+
+from __future__ import annotations
+
+import logging
+from typing import Self
+
+from roborock.data import NetworkInfo
+from roborock.devices.cache import Cache
+from roborock.devices.traits.v1 import common
+from roborock.roborock_typing import RoborockCommand
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class NetworkInfoTrait(NetworkInfo, common.V1TraitMixin):
+    """Trait for device network information.
+
+    This trait will always prefer reading from the cache if available. This
+    information is usually already fetched when creating the device local
+    connection, so reading from the cache avoids an unnecessary RPC call.
+    However, we have the fallback to reading from the device if the cache is
+    not populated for some reason.
+    """
+
+    command = RoborockCommand.GET_NETWORK_INFO
+
+    def __init__(self, device_uid: str, cache: Cache) -> None:
+        """Initialize the trait."""
+        self._device_uid = device_uid
+        self._cache = cache
+
+    async def refresh(self) -> Self:
+        """Refresh the network info from the cache."""
+
+        cache_data = await self._cache.get()
+        if cache_data.network_info and (network_info := cache_data.network_info.get(self._device_uid)):
+            _LOGGER.debug("Using cached network info for device %s", self._device_uid)
+            self._update_trait_values(network_info)
+            return self
+
+        # Load from device if not in cache
+        _LOGGER.debug("No cached network info for device %s, fetching from device", self._device_uid)
+        await super().refresh()
+
+        # Update the cache with the new network info
+        cache_data.network_info[self._device_uid] = self
+        await self._cache.set(cache_data)
+
+        return self
+
+    def _parse_response(self, response: common.V1ResponseData) -> NetworkInfo:
+        """Parse the response from the device into a NetworkInfo."""
+        if not isinstance(response, dict):
+            raise ValueError(f"Unexpected NetworkInfoTrait response format: {response!r}")
+        return NetworkInfo.from_dict(response)

tests/devices/test_v1_device.py

@@ -55,7 +55,15 @@ def device_fixture(channel: AsyncMock, rpc_channel: AsyncMock, mqtt_rpc_channel:
         device_info=HOME_DATA.devices[0],
         product=HOME_DATA.products[0],
         channel=channel,
-        trait=v1.create(HOME_DATA.products[0], HOME_DATA, rpc_channel, mqtt_rpc_channel, AsyncMock(), NoCache()),
+        trait=v1.create(
+            HOME_DATA.devices[0].duid,
+            HOME_DATA.products[0],
+            HOME_DATA,
+            rpc_channel,
+            mqtt_rpc_channel,
+            AsyncMock(),
+            NoCache(),
+        ),
     )
 
 

tests/devices/traits/v1/fixtures.py

@@ -60,6 +60,7 @@ def device_fixture(
         product=HOME_DATA.products[0],
         channel=channel,
         trait=v1.create(
+            HOME_DATA.devices[0].duid,
             HOME_DATA.products[0],
             HOME_DATA,
             mock_rpc_channel,

tests/devices/traits/v1/test_network_info.py

@@ -0,0 +1,63 @@
+"""Tests for the NetworkInfoTrait class."""
+
+from unittest.mock import AsyncMock
+
+import pytest
+
+from roborock.data import NetworkInfo
+from roborock.devices.cache import Cache
+from roborock.devices.device import RoborockDevice
+from roborock.devices.traits.v1.network_info import NetworkInfoTrait
+from roborock.roborock_typing import RoborockCommand
+from tests.mock_data import NETWORK_INFO
+
+DEVICE_UID = "abc123"
+
+
+@pytest.fixture
+def network_info_trait(device: RoborockDevice) -> NetworkInfoTrait:
+    """Create a NetworkInfoTrait instance with mocked dependencies."""
+    assert device.v1_properties
+    return device.v1_properties.network_info
+
+
+async def test_network_info_from_cache(
+    network_info_trait: NetworkInfoTrait, roborock_cache: Cache, mock_rpc_channel: AsyncMock
+) -> None:
+    """Test that network info is read from the cache."""
+    cache_data = await roborock_cache.get()
+    network_info = NetworkInfo.from_dict(NETWORK_INFO)
+    cache_data.network_info[DEVICE_UID] = network_info
+    await roborock_cache.set(cache_data)
+
+    await network_info_trait.refresh()
+
+    assert network_info_trait.ip == "1.1.1.1"
+    assert network_info_trait.mac == "aa:bb:cc:dd:ee:ff"
+    assert network_info_trait.bssid == "aa:bb:cc:dd:ee:ff"
+    assert network_info_trait.rssi == -50
+    mock_rpc_channel.send_command.assert_not_called()
+
+
+async def test_network_info_from_device(
+    network_info_trait: NetworkInfoTrait, roborock_cache: Cache, mock_rpc_channel: AsyncMock
+) -> None:
+    """Test that network info is fetched from the device when not in cache."""
+    mock_rpc_channel.send_command.return_value = {
+        **NETWORK_INFO,
+        "ip": "2.2.2.2",
+    }
+
+    await network_info_trait.refresh()
+
+    assert network_info_trait.ip == "2.2.2.2"
+    assert network_info_trait.mac == "aa:bb:cc:dd:ee:ff"
+    assert network_info_trait.bssid == "aa:bb:cc:dd:ee:ff"
+    assert network_info_trait.rssi == -50
+    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_NETWORK_INFO)
+
+    # Verify it's now in the cache
+    cache_data = await roborock_cache.get()
+    cached_info = cache_data.network_info.get(DEVICE_UID)
+    assert cached_info
+    assert cached_info.ip == "2.2.2.2"