feat: Add a DeviceManager to perform discovery (#399)

* feat: Add a DeviceManager to perform discovery

* feat: Update review feedback

* feat: Update tests with additional feedback

* feat: Fix lint error

---------

Co-authored-by: semantic-release <semantic-release>

Author: allenporter

roborock/containers.py

@@ -7,6 +7,7 @@
 from dataclasses import asdict, dataclass, field
 from datetime import timezone
 from enum import Enum
+from functools import cached_property
 from typing import Any, NamedTuple, get_args, get_origin
 
 from .code_mappings import (
@@ -469,6 +470,21 @@ def get_all_devices(self) -> list[HomeDataDevice]:
             devices += self.received_devices
         return devices
 
+    @cached_property
+    def product_map(self) -> dict[str, HomeDataProduct]:
+        """Returns a dictionary of product IDs to HomeDataProduct objects."""
+        return {product.id: product for product in self.products}
+
+    @cached_property
+    def device_products(self) -> dict[str, tuple[HomeDataDevice, HomeDataProduct]]:
+        """Returns a dictionary of device DUIDs to HomeDataDeviceProduct objects."""
+        product_map = self.product_map
+        return {
+            device.duid: (device, product)
+            for device in self.get_all_devices()
+            if (product := product_map.get(device.product_id)) is not None
+        }
+
 
 @dataclass
 class LoginData(RoborockBase):

roborock/devices/README.md

@@ -0,0 +1,98 @@
+# Roborock Device Discovery
+
+This page documents the full lifecycle of device discovery across Cloud and Network.
+
+## Init account setup
+
+### Login
+
+- Login can happen with either email and password or email and sending a code. We
+  currently prefer email with sending a code -- however the roborock no longer
+  supports this method of login. In the future we may want to migrate to password
+  if this login method is no longer supported.
+- The Login API provides a `userData` object with information on connecting to the cloud APIs
+- This `rriot` data contains per-session information, unique each time you login.
+  - This contains information used to connect to MQTT
+  - You get an `-eu` suffix in the API URLs if you are in the eu and `-us` if you are in the us
+
+## Home Data
+
+The `HomeData` includes information about the various devices in the home. We use `v3`
+and it is notable that if devices don't show up in the `home_data` response it is likely
+that a newer version of the API should be used.
+
+- `products`: This is a list of all of the products you have on your account. These objects are always the same (i.e. a s7 maxv is always the exact same.)
+  - It only shows the products for devices available on your account
+- `devices` and `received_devices`:
+  - These both share the same objects, but one is for devices that have been shared with you and one is those that are on your account.
+  - The big things here are (MOST are static):
+    - `duid`: A unique identifier for your device (this is always the same i think)
+    - `name`: The name of the device in your app
+    - `local_key`: The local key that is needed for encoding and decoding messages for the device. This stays the same unless someone sets their vacuum back up.
+    - `pv`: the protocol version (i.e. 1.0 or A1 or B1)
+    - `product_id`: The id of the product from the above products list.
+    - `device_status`: An initial status for some of the data we care about, though this changes on each update.
+- `rooms`: The rooms in the home.
+  - This changes if the user adds a new room or changes its name.
+  - We have to combine this with the room numbers from `GET_ROOM_MAPPING` on the device
+  - There is another REST request `get_rooms` that will do the same thing.
+  - Note: If we cache home_data, we likely need to use `get_rooms` to get rooms fresh
+
+## Device Connections
+
+### MQTT connection
+
+- Initial device information must be obtained from MQTT
+- We typically set up the MQTT device connection before the local device connection.
+  - The `NetworkingInfo` needs to be fetched to get additional information about connecting to the device:
+    - e.g. Local IP Address
+  - This networking info can be cached to reduce network calls
+  - MQTT also is the only way to get the device Map
+- Incoming and outgoing messages are decoded/encoded using the device `local_key`
+- Otherwise all commands may be performed locally.
+
+## Local connection
+
+- We can use the `ip` from the `NetworkingInfo` to find the device
+- The local connection is preferred to for improved latency and reducing load on the cloud servers to avoid rate limiting.
+- Connections are made using a normal TCP socket on port `58867`
+- Incoming and outgoing messages are decoded/encoded using the device `local_key`
+- Messages received on the stream may be partially received so we keep a running as messages are partially decoded
+
+## Design
+
+### Current API Issues
+
+- Complex Inheritance Hierarchy: Multiple inheritance with classes like RoborockMqttClientV1 inheriting from both RoborockMqttClient and RoborockClientV1
+
+- Callback-Heavy Design: Heavy reliance on callbacks and listeners in RoborockClientV1.on_message_received and the ListenerModel system
+
+- Version Fragmentation: Separate v1 and A01 APIs with different patterns and abstractions
+
+- Mixed Concerns: Classes handle both communication protocols (MQTT/local) and device-specific logic
+
+- Complex Caching: The AttributeCache system with RepeatableTask adds complexity
+
+- Manual Connection Management: Users need to manually set up both MQTT and local clients as shown in the README example
+
+## Design Changes
+
+- Prefer a single unfieid client that handles both MQTT and local connections internally.
+
+- Home and device discovery (fetching home data and device setup) will be behind a single API.
+
+- Asyncio First: Everything should be asyncio as much as possible, with fewer callbacks.
+
+- The clients should be working in terms of devices. We need to detect capabilities for each device and not expose details about API versions.
+
+- Reliability issues: The current Home Assistant integration has issues with reliability and needs to be simplified. It may be that there are bugs with the exception handling and it's too heavy the cloud APIs and could benefit from more seamless caching.
+
+## Implementation Details
+
+- We don't really need to worry about backwards compatibility for the new set of APIs.
+
+- We'll have a `RoborockManager` responsible for managing the connections and getting devices.
+
+- Caching can be persisted to disk. The caller can implement the cache storage themselves, but we need to give them an API to do so.
+
+- Users don't really choose between cloud vs local. However, we will want to allow the caller to know if its using the locale connection so we can show a warnings.

roborock/devices/__init__.py

@@ -0,0 +1,6 @@
+"""The devices module provides functionality to discover Roborock devices on the network."""
+
+__all__ = [
+    "device",
+    "device_manager",
+]

roborock/devices/device.py

@@ -0,0 +1,65 @@
+"""Module for Roborock devices.
+
+This interface is experimental and subject to breaking changes without notice
+until the API is stable.
+"""
+
+import enum
+import logging
+from functools import cached_property
+
+from roborock.containers import HomeDataDevice, HomeDataProduct, UserData
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "RoborockDevice",
+    "DeviceVersion",
+]
+
+
+class DeviceVersion(enum.StrEnum):
+    """Enum for device versions."""
+
+    V1 = "1.0"
+    A01 = "A01"
+    UNKNOWN = "unknown"
+
+
+class RoborockDevice:
+    """Unified Roborock device class with automatic connection setup."""
+
+    def __init__(self, user_data: UserData, device_info: HomeDataDevice, product_info: HomeDataProduct) -> None:
+        """Initialize the RoborockDevice with device info, user data, and capabilities."""
+        self._user_data = user_data
+        self._device_info = device_info
+        self._product_info = product_info
+
+    @property
+    def duid(self) -> str:
+        """Return the device unique identifier (DUID)."""
+        return self._device_info.duid
+
+    @property
+    def name(self) -> str:
+        """Return the device name."""
+        return self._device_info.name
+
+    @cached_property
+    def device_version(self) -> str:
+        """Return the device version.
+
+        At the moment this is a simple check against the product version (pv) of the device
+        and used as a placeholder for upcoming functionality for devices that will behave
+        differently based on the version and capabilities.
+        """
+        if self._device_info.pv == DeviceVersion.V1.value:
+            return DeviceVersion.V1
+        elif self._device_info.pv == DeviceVersion.A01.value:
+            return DeviceVersion.A01
+        _LOGGER.warning(
+            "Unknown device version %s for device %s, using default UNKNOWN",
+            self._device_info.pv,
+            self._device_info.name,
+        )
+        return DeviceVersion.UNKNOWN

roborock/devices/device_manager.py

@@ -0,0 +1,91 @@
+"""Module for discovering Roborock devices."""
+
+import logging
+from collections.abc import Awaitable, Callable
+
+from roborock.containers import (
+    HomeData,
+    HomeDataDevice,
+    HomeDataProduct,
+    UserData,
+)
+from roborock.devices.device import RoborockDevice
+from roborock.web_api import RoborockApiClient
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "create_device_manager",
+    "create_home_data_api",
+    "DeviceManager",
+    "HomeDataApi",
+    "DeviceCreator",
+]
+
+
+HomeDataApi = Callable[[], Awaitable[HomeData]]
+DeviceCreator = Callable[[HomeDataDevice, HomeDataProduct], RoborockDevice]
+
+
+class DeviceManager:
+    """Central manager for Roborock device discovery and connections."""
+
+    def __init__(
+        self,
+        home_data_api: HomeDataApi,
+        device_creator: DeviceCreator,
+    ) -> None:
+        """Initialize the DeviceManager with user data and optional cache storage."""
+        self._home_data_api = home_data_api
+        self._device_creator = device_creator
+        self._devices: dict[str, RoborockDevice] = {}
+
+    async def discover_devices(self) -> list[RoborockDevice]:
+        """Discover all devices for the logged-in user."""
+        home_data = await self._home_data_api()
+        device_products = home_data.device_products
+        _LOGGER.debug("Discovered %d devices %s", len(device_products), home_data)
+
+        self._devices = {
+            duid: self._device_creator(device, product) for duid, (device, product) in device_products.items()
+        }
+        return list(self._devices.values())
+
+    async def get_device(self, duid: str) -> RoborockDevice | None:
+        """Get a specific device by DUID."""
+        return self._devices.get(duid)
+
+    async def get_devices(self) -> list[RoborockDevice]:
+        """Get all discovered devices."""
+        return list(self._devices.values())
+
+
+def create_home_data_api(email: str, user_data: UserData) -> HomeDataApi:
+    """Create a home data API wrapper.
+
+    This function creates a wrapper around the Roborock API client to fetch
+    home data for the user.
+    """
+
+    client = RoborockApiClient(email, user_data)
+
+    async def home_data_api() -> HomeData:
+        return await client.get_home_data(user_data)
+
+    return home_data_api
+
+
+async def create_device_manager(user_data: UserData, home_data_api: HomeDataApi) -> 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.
+    """
+
+    def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
+        return RoborockDevice(user_data, device, product)
+
+    manager = DeviceManager(home_data_api, device_creator)
+    await manager.discover_devices()
+    return manager

roborock/mqtt/__init__.py

@@ -4,4 +4,7 @@
 modules.
 """
 
+# This module is part of the Roborock Python library, which provides a way to
+# interact with Roborock devices using MQTT. It is not intended to be used directly,
+# but rather as a base for higher level modules.
 __all__: list[str] = []

tests/devices/__init__.py

@@ -0,0 +1 @@
+"""Tests for the device module."""

tests/devices/test_device_manager.py

@@ -0,0 +1,81 @@
+"""Tests for the DeviceManager class."""
+
+from unittest.mock import patch
+
+import pytest
+
+from roborock.containers import HomeData, UserData
+from roborock.devices.device import DeviceVersion
+from roborock.devices.device_manager import create_device_manager, create_home_data_api
+from roborock.exceptions import RoborockException
+
+from .. import mock_data
+
+USER_DATA = UserData.from_dict(mock_data.USER_DATA)
+
+
+async def home_home_data_no_devices() -> HomeData:
+    """Mock home data API that returns no devices."""
+    return HomeData(
+        id=1,
+        name="Test Home",
+        devices=[],
+        products=[],
+    )
+
+
+async def mock_home_data() -> HomeData:
+    """Mock home data API that returns devices."""
+    return HomeData.from_dict(mock_data.HOME_DATA_RAW)
+
+
+async def test_no_devices() -> None:
+    """Test the DeviceManager created with no devices returned from the API."""
+
+    device_manager = await create_device_manager(USER_DATA, home_home_data_no_devices)
+    devices = await device_manager.get_devices()
+    assert devices == []
+
+
+async def test_with_device() -> None:
+    """Test the DeviceManager created with devices returned from the API."""
+    device_manager = await create_device_manager(USER_DATA, mock_home_data)
+    devices = await device_manager.get_devices()
+    assert len(devices) == 1
+    assert devices[0].duid == "abc123"
+    assert devices[0].name == "Roborock S7 MaxV"
+    assert devices[0].device_version == DeviceVersion.V1
+
+    device = await device_manager.get_device("abc123")
+    assert device is not None
+    assert device.duid == "abc123"
+    assert device.name == "Roborock S7 MaxV"
+    assert device.device_version == DeviceVersion.V1
+
+
+async def test_get_non_existent_device() -> None:
+    """Test getting a non-existent device."""
+    device_manager = await create_device_manager(USER_DATA, mock_home_data)
+    device = await device_manager.get_device("non_existent_duid")
+    assert device is None
+
+
+async def test_home_data_api_exception() -> None:
+    """Test the home data API with an exception."""
+
+    async def home_data_api_exception() -> HomeData:
+        raise RoborockException("Test exception")
+
+    with pytest.raises(RoborockException, match="Test exception"):
+        await create_device_manager(USER_DATA, home_data_api_exception)
+
+
+async def test_create_home_data_api_exception() -> None:
+    """Test that exceptions from the home data API are propagated through the wrapper."""
+
+    with patch("roborock.devices.device_manager.RoborockApiClient.get_home_data") as mock_get_home_data:
+        mock_get_home_data.side_effect = RoborockException("Test exception")
+        api = create_home_data_api(USER_DATA, mock_get_home_data)
+
+        with pytest.raises(RoborockException, match="Test exception"):
+            await api()

tests/mock_data.py

@@ -293,7 +293,7 @@
             "runtimeEnv": None,
             "timeZoneId": "America/Los_Angeles",
             "iconUrl": "no_url",
-            "productId": "product123",
+            "productId": PRODUCT_ID,
             "lon": None,
             "lat": None,
             "share": False,