From 914d8973d90d3cb27f3274bf811fc8a9553cdf23 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Thu, 23 Oct 2025 07:47:00 -0700 Subject: [PATCH 1/2] 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. --- roborock/cli.py | 11 ++++ roborock/devices/device_manager.py | 1 + roborock/devices/traits/v1/__init__.py | 18 +++++- roborock/devices/traits/v1/network_info.py | 56 +++++++++++++++++ tests/devices/test_v1_device.py | 10 +++- tests/devices/traits/v1/fixtures.py | 1 + tests/devices/traits/v1/test_network_info.py | 63 ++++++++++++++++++++ 7 files changed, 158 insertions(+), 2 deletions(-) create mode 100644 roborock/devices/traits/v1/network_info.py create mode 100644 tests/devices/traits/v1/test_network_info.py diff --git a/roborock/cli.py b/roborock/cli.py index 8829137d..a237c266 100644 --- a/roborock/cli.py +++ b/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(): diff --git a/roborock/devices/device_manager.py b/roborock/devices/device_manager.py index ee6e23ba..5218a839 100644 --- a/roborock/devices/device_manager.py +++ b/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, diff --git a/roborock/devices/traits/v1/__init__.py b/roborock/devices/traits/v1/__init__.py index aabc28f9..0b128e79 100644 --- a/roborock/devices/traits/v1/__init__.py +++ b/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, + ) diff --git a/roborock/devices/traits/v1/network_info.py b/roborock/devices/traits/v1/network_info.py new file mode 100644 index 00000000..e11cd778 --- /dev/null +++ b/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) diff --git a/tests/devices/test_v1_device.py b/tests/devices/test_v1_device.py index ff1f613c..78b54607 100644 --- a/tests/devices/test_v1_device.py +++ b/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(), + ), ) diff --git a/tests/devices/traits/v1/fixtures.py b/tests/devices/traits/v1/fixtures.py index 1026ea3e..01aca306 100644 --- a/tests/devices/traits/v1/fixtures.py +++ b/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, diff --git a/tests/devices/traits/v1/test_network_info.py b/tests/devices/traits/v1/test_network_info.py new file mode 100644 index 00000000..da0e5486 --- /dev/null +++ b/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() + assert "roborock.vacuum.a15" in cache_data.network_info + cached_info = cache_data.network_info[DEVICE_UID] + assert cached_info.ip == "2.2.2.2" From 4bab32b23d4eaf58476e5e462480d843a390d77c Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Thu, 23 Oct 2025 07:48:55 -0700 Subject: [PATCH 2/2] chore: Update test assertion for network info --- tests/devices/traits/v1/test_network_info.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/devices/traits/v1/test_network_info.py b/tests/devices/traits/v1/test_network_info.py index da0e5486..2da1105e 100644 --- a/tests/devices/traits/v1/test_network_info.py +++ b/tests/devices/traits/v1/test_network_info.py @@ -58,6 +58,6 @@ async def test_network_info_from_device( # Verify it's now in the cache cache_data = await roborock_cache.get() - assert "roborock.vacuum.a15" in cache_data.network_info - cached_info = cache_data.network_info[DEVICE_UID] + cached_info = cache_data.network_info.get(DEVICE_UID) + assert cached_info assert cached_info.ip == "2.2.2.2"