Merge branch 'main' into supported_features_markdown

Author: Lash-L

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "python-roborock"
-version = "2.31.0"
+version = "2.33.0"
 description = "A package to control Roborock vacuums."
 authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
 license = "GPL-3.0-only"
@@ -39,7 +39,7 @@ requires = ["poetry-core==1.8.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.group.dev.dependencies]
-pytest-asyncio = "*"
+pytest-asyncio = ">=1.1.0"
 pytest = "*"
 pre-commit = ">=3.5,<5.0"
 mypy = "*"

pyproject.toml

@@ -12,6 +12,7 @@
 
 from roborock import SHORT_MODEL_TO_ENUM, DeviceFeatures, RoborockCommand, RoborockException
 from roborock.containers import DeviceData, HomeData, HomeDataProduct, LoginData, NetworkInfo, RoborockBase, UserData
+from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device_manager import create_device_manager, create_home_data_api
 from roborock.protocol import MessageParser
 from roborock.util import run_sync
@@ -39,7 +40,7 @@ class ConnectionCache(RoborockBase):
     network_info: dict[str, NetworkInfo] | None = None
 
 
-class RoborockContext:
+class RoborockContext(Cache):
     roborock_file = Path("~/.roborock").expanduser()
     _cache_data: ConnectionCache | None = None
 
@@ -68,6 +69,18 @@ def cache_data(self) -> ConnectionCache:
         self.validate()
         return self._cache_data
 
+    async def get(self) -> CacheData:
+        """Get cached value."""
+        connection_cache = self.cache_data()
+        return CacheData(home_data=connection_cache.home_data, network_info=connection_cache.network_info or {})
+
+    async def set(self, value: CacheData) -> None:
+        """Set value in the cache."""
+        connection_cache = self.cache_data()
+        connection_cache.home_data = value.home_data
+        connection_cache.network_info = value.network_info
+        self.update(connection_cache)
+
 
 @click.option("-d", "--debug", default=False, count=True)
 @click.version_option(package_name="python-roborock")
@@ -119,14 +132,8 @@ async def session(ctx, duration: int):
 
     home_data_api = create_home_data_api(cache_data.email, cache_data.user_data)
 
-    async def home_data_cache() -> HomeData:
-        if cache_data.home_data is None:
-            cache_data.home_data = await home_data_api()
-            context.update(cache_data)
-        return cache_data.home_data
-
     # Create device manager
-    device_manager = await create_device_manager(cache_data.user_data, home_data_cache)
+    device_manager = await create_device_manager(cache_data.user_data, home_data_api, context)
 
     devices = await device_manager.get_devices()
     click.echo(f"Discovered devices: {', '.join([device.name for device in devices])}")

roborock/cli.py

@@ -425,6 +425,13 @@ def get_mop_mode_code(self, mop_mode: str) -> int:
             raise RoborockException("Attempted to get mop_mode before status has been updated.")
         return self.mop_mode.as_dict().get(mop_mode)
 
+    @property
+    def current_map(self) -> int | None:
+        """Returns the current map ID if the map is present."""
+        if self.map_status is not None:
+            return (self.map_status - 3) // 4
+        return None
+
 
 @dataclass
 class S4MaxStatus(Status):

roborock/containers.py

@@ -3,4 +3,5 @@
 __all__ = [
     "device",
     "device_manager",
+    "cache",
 ]

roborock/devices/__init__.py

@@ -0,0 +1,57 @@
+"""This module provides caching functionality for the Roborock device management system.
+
+This module defines a cache interface that you may use to cache device
+information to avoid unnecessary API calls. Callers may implement
+this interface to provide their own caching mechanism.
+"""
+
+from dataclasses import dataclass, field
+from typing import Protocol
+
+from roborock.containers import HomeData, NetworkInfo
+
+
+@dataclass
+class CacheData:
+    """Data structure for caching device information."""
+
+    home_data: HomeData | None = None
+    """Home data containing device and product information."""
+
+    network_info: dict[str, NetworkInfo] = field(default_factory=dict)
+    """Network information indexed by device DUID."""
+
+
+class Cache(Protocol):
+    """Protocol for a cache that can store and retrieve values."""
+
+    async def get(self) -> CacheData:
+        """Get cached value."""
+        ...
+
+    async def set(self, value: CacheData) -> None:
+        """Set value in the cache."""
+        ...
+
+
+class InMemoryCache(Cache):
+    """In-memory cache implementation."""
+
+    def __init__(self):
+        self._data = CacheData()
+
+    async def get(self) -> CacheData:
+        return self._data
+
+    async def set(self, value: CacheData) -> None:
+        self._data = value
+
+
+class NoCache(Cache):
+    """No-op cache implementation."""
+
+    async def get(self) -> CacheData:
+        return CacheData()
+
+    async def set(self, value: CacheData) -> None:
+        pass

roborock/devices/cache.py

@@ -18,6 +18,7 @@
 from roborock.protocol import create_mqtt_params
 from roborock.web_api import RoborockApiClient
 
+from .cache import Cache, NoCache
 from .channel import Channel
 from .mqtt_channel import create_mqtt_channel
 from .traits.dyad import DyadApi
@@ -32,8 +33,6 @@
     "create_device_manager",
     "create_home_data_api",
     "DeviceManager",
-    "HomeDataApi",
-    "DeviceCreator",
 ]
 
 
@@ -57,19 +56,27 @@ def __init__(
         home_data_api: HomeDataApi,
         device_creator: DeviceCreator,
         mqtt_session: MqttSession,
+        cache: Cache,
     ) -> None:
         """Initialize the DeviceManager with user data and optional cache storage.
 
         This takes ownership of the MQTT session and will close it when the manager is closed.
         """
         self._home_data_api = home_data_api
+        self._cache = cache
         self._device_creator = device_creator
         self._devices: dict[str, RoborockDevice] = {}
         self._mqtt_session = mqtt_session
 
     async def discover_devices(self) -> list[RoborockDevice]:
         """Discover all devices for the logged-in user."""
-        home_data = await self._home_data_api()
+        cache_data = await self._cache.get()
+        if not cache_data.home_data:
+            _LOGGER.debug("No cached home data found, fetching from API")
+            cache_data.home_data = await self._home_data_api()
+            await self._cache.set(cache_data)
+        home_data = cache_data.home_data
+
         device_products = home_data.device_products
         _LOGGER.debug("Discovered %d devices %s", len(device_products), home_data)
 
@@ -118,13 +125,19 @@ async def home_data_api() -> HomeData:
     return home_data_api
 
 
-async def create_device_manager(user_data: UserData, home_data_api: HomeDataApi) -> DeviceManager:
+async def create_device_manager(
+    user_data: UserData,
+    home_data_api: HomeDataApi,
+    cache: Cache | None = None,
+) -> DeviceManager:
     """Convenience function to create and initialize a DeviceManager.
 
     The Home Data is fetched using the provided home_data_api callable which
     is exposed this way to allow for swapping out other implementations to
     include caching or other optimizations.
     """
+    if cache is None:
+        cache = NoCache()
 
     mqtt_params = create_mqtt_params(user_data.rriot)
     mqtt_session = await create_mqtt_session(mqtt_params)
@@ -135,7 +148,7 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
         # TODO: Define a registration mechanism/factory for v1 traits
         match device.pv:
             case DeviceVersion.V1:
-                channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device)
+                channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, cache)
                 traits.append(StatusTrait(product, channel.rpc_channel))
             case DeviceVersion.A01:
                 mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
@@ -150,6 +163,6 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
                 raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}")
         return RoborockDevice(device, channel, traits)
 
-    manager = DeviceManager(home_data_api, device_creator, mqtt_session=mqtt_session)
+    manager = DeviceManager(home_data_api, device_creator, mqtt_session=mqtt_session, cache=cache)
     await manager.discover_devices()
     return manager

roborock/devices/device_manager.py

@@ -18,6 +18,7 @@
 from roborock.roborock_message import RoborockMessage
 from roborock.roborock_typing import RoborockCommand
 
+from .cache import Cache
 from .channel import Channel
 from .local_channel import LocalChannel, LocalSession, create_local_session
 from .mqtt_channel import MqttChannel
@@ -46,6 +47,7 @@ def __init__(
         security_data: SecurityData,
         mqtt_channel: MqttChannel,
         local_session: LocalSession,
+        cache: Cache,
     ) -> None:
         """Initialize the V1Channel.
 
@@ -62,7 +64,7 @@ def __init__(
         self._mqtt_unsub: Callable[[], None] | None = None
         self._local_unsub: Callable[[], None] | None = None
         self._callback: Callable[[RoborockMessage], None] | None = None
-        self._networking_info: NetworkInfo | None = None
+        self._cache = cache
 
     @property
     def is_connected(self) -> bool:
@@ -131,19 +133,26 @@ async def _get_networking_info(self) -> NetworkInfo:
 
         This is a cloud only command used to get the local device's IP address.
         """
+        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)
+            return network_info
         try:
-            return await self._mqtt_rpc_channel.send_command(
+            network_info = await self._mqtt_rpc_channel.send_command(
                 RoborockCommand.GET_NETWORK_INFO, response_type=NetworkInfo
             )
         except RoborockException as e:
             raise RoborockException(f"Network info failed for device {self._device_uid}") from e
+        _LOGGER.debug("Network info for device %s: %s", self._device_uid, network_info)
+        cache_data.network_info[self._device_uid] = network_info
+        await self._cache.set(cache_data)
+        return network_info
 
     async def _local_connect(self) -> Callable[[], None]:
         """Set up local connection if possible."""
         _LOGGER.debug("Attempting to connect to local channel for device %s", self._device_uid)
-        if self._networking_info is None:
-            self._networking_info = await self._get_networking_info()
-        host = self._networking_info.ip
+        networking_info = await self._get_networking_info()
+        host = networking_info.ip
         _LOGGER.debug("Connecting to local channel at %s", host)
         self._local_channel = self._local_session(host)
         try:
@@ -168,10 +177,14 @@ def _on_local_message(self, message: RoborockMessage) -> None:
 
 
 def create_v1_channel(
-    user_data: UserData, mqtt_params: MqttParams, mqtt_session: MqttSession, device: HomeDataDevice
+    user_data: UserData,
+    mqtt_params: MqttParams,
+    mqtt_session: MqttSession,
+    device: HomeDataDevice,
+    cache: Cache,
 ) -> V1Channel:
     """Create a V1Channel for the given device."""
     security_data = create_security_data(user_data.rriot)
     mqtt_channel = MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)
     local_session = create_local_session(device.local_key)
-    return V1Channel(device.duid, security_data, mqtt_channel, local_session=local_session)
+    return V1Channel(device.duid, security_data, mqtt_channel, local_session=local_session, cache=cache)

roborock/devices/v1_channel.py

@@ -125,14 +125,14 @@ async def _run_task(self, start_future: asyncio.Future[None] | None) -> None:
             except Exception as err:
                 # This error is thrown when the MQTT loop is cancelled
                 # and the generator is not stopped.
-                if "generator didn't stop" in str(err):
+                if "generator didn't stop" in str(err) or "generator didn't yield" in str(err):
                     _LOGGER.debug("MQTT loop was cancelled")
                     return
                 if start_future:
                     _LOGGER.error("Uncaught error starting MQTT session: %s", err)
                     start_future.set_exception(err)
                     return
-                _LOGGER.error("Uncaught error during MQTT session: %s", err)
+                _LOGGER.exception("Uncaught error during MQTT session: %s", err)
 
             self._healthy = False
             _LOGGER.info("MQTT session disconnected, retrying in %s seconds", self._backoff.total_seconds())
@@ -180,7 +180,7 @@ async def _process_message_loop(self, client: aiomqtt.Client) -> None:
                 except asyncio.CancelledError:
                     raise
                 except Exception as e:
-                    _LOGGER.error("Uncaught exception in subscriber callback: %s", e)
+                    _LOGGER.exception("Uncaught exception in subscriber callback: %s", e)
 
     async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
         """Subscribe to messages on the specified topic and invoke the callback for new messages.

roborock/mqtt/roborock_session.py

@@ -6,6 +6,7 @@
 import pytest
 
 from roborock.containers import HomeData, UserData
+from roborock.devices.cache import CacheData, InMemoryCache
 from roborock.devices.device_manager import create_device_manager, create_home_data_api
 from roborock.exceptions import RoborockException
 
@@ -98,3 +99,37 @@ async def test_create_home_data_api_exception() -> None:
 
         with pytest.raises(RoborockException, match="Test exception"):
             await api()
+
+
+async def test_cache_logic() -> None:
+    """Test that the cache logic works correctly."""
+    call_count = 0
+
+    async def mock_home_data_with_counter() -> HomeData:
+        nonlocal call_count
+        call_count += 1
+        return HomeData.from_dict(mock_data.HOME_DATA_RAW)
+
+    class TestCache:
+        def __init__(self):
+            self._data = CacheData()
+
+        async def get(self) -> CacheData:
+            return self._data
+
+        async def set(self, value: CacheData) -> None:
+            self._data = value
+
+    # First call happens during create_device_manager initialization
+    device_manager = await create_device_manager(USER_DATA, mock_home_data_with_counter, cache=InMemoryCache())
+    assert call_count == 1
+
+    # Second call should use cache, not increment call_count
+    devices2 = await device_manager.discover_devices()
+    assert call_count == 1  # Should still be 1, not 2
+    assert len(devices2) == 1
+
+    await device_manager.close()
+    assert len(devices2) == 1
+
+    await device_manager.close()