roborock.devices.device_manager

  1"""Module for discovering Roborock devices."""
  2
  3import asyncio
  4import enum
  5import logging
  6from collections.abc import Callable, Mapping
  7from dataclasses import dataclass
  8from typing import Any
  9
 10import aiohttp
 11
 12from roborock.data import (
 13    HomeData,
 14    HomeDataDevice,
 15    HomeDataProduct,
 16    UserData,
 17)
 18from roborock.devices.device import DeviceReadyCallback, RoborockDevice
 19from roborock.diagnostics import Diagnostics, redact_device_data
 20from roborock.exceptions import RoborockException
 21from roborock.map.map_parser import MapParserConfig
 22from roborock.mqtt.roborock_session import create_lazy_mqtt_session
 23from roborock.mqtt.session import MqttSession, SessionUnauthorizedHook
 24from roborock.protocol import create_mqtt_params
 25from roborock.web_api import RoborockApiClient, UserWebApiClient
 26
 27from .cache import Cache, DeviceCache, NoCache
 28from .rpc.v1_channel import create_v1_channel
 29from .traits import Trait, a01, b01, v1
 30from .transport.channel import Channel
 31from .transport.mqtt_channel import create_mqtt_channel
 32
 33_LOGGER = logging.getLogger(__name__)
 34
 35__all__ = [
 36    "create_device_manager",
 37    "UserParams",
 38    "DeviceManager",
 39]
 40
 41
 42DeviceCreator = Callable[[HomeData, HomeDataDevice, HomeDataProduct], RoborockDevice]
 43
 44
 45class DeviceVersion(enum.StrEnum):
 46    """Enum for device versions."""
 47
 48    V1 = "1.0"
 49    A01 = "A01"
 50    B01 = "B01"
 51    UNKNOWN = "unknown"
 52
 53
 54class UnsupportedDeviceError(RoborockException):
 55    """Exception raised when a device is unsupported."""
 56
 57
 58class DeviceManager:
 59    """Central manager for Roborock device discovery and connections."""
 60
 61    def __init__(
 62        self,
 63        web_api: UserWebApiClient,
 64        device_creator: DeviceCreator,
 65        mqtt_session: MqttSession,
 66        cache: Cache,
 67        diagnostics: Diagnostics,
 68    ) -> None:
 69        """Initialize the DeviceManager with user data and optional cache storage.
 70
 71        This takes ownership of the MQTT session and will close it when the manager is closed.
 72        """
 73        self._web_api = web_api
 74        self._cache = cache
 75        self._device_creator = device_creator
 76        self._devices: dict[str, RoborockDevice] = {}
 77        self._mqtt_session = mqtt_session
 78        self._diagnostics = diagnostics
 79        self._home_data: HomeData | None = None
 80
 81    async def discover_devices(self, prefer_cache: bool = True) -> list[RoborockDevice]:
 82        """Discover all devices for the logged-in user."""
 83        self._diagnostics.increment("discover_devices")
 84        cache_data = await self._cache.get()
 85        if not cache_data.home_data or not prefer_cache:
 86            _LOGGER.debug("Fetching home data (prefer_cache=%s)", prefer_cache)
 87            self._diagnostics.increment("fetch_home_data")
 88            try:
 89                cache_data.home_data = await self._web_api.get_home_data()
 90            except RoborockException as ex:
 91                if not cache_data.home_data:
 92                    raise
 93                _LOGGER.debug("Failed to fetch home data, using cached data: %s", ex)
 94            await self._cache.set(cache_data)
 95        self._home_data = cache_data.home_data
 96
 97        device_products = self._home_data.device_products
 98        _LOGGER.debug("Discovered %d devices", len(device_products))
 99
100        # These are connected serially to avoid overwhelming the MQTT broker
101        new_devices = {}
102        start_tasks = []
103        supported_devices_counter = self._diagnostics.subkey("supported_devices")
104        unsupported_devices_counter = self._diagnostics.subkey("unsupported_devices")
105        for duid, (device, product) in device_products.items():
106            _LOGGER.debug("[%s] Discovered device %s %s", duid, product.summary_info(), device.summary_info())
107            if duid in self._devices:
108                continue
109            try:
110                new_device = self._device_creator(self._home_data, device, product)
111            except UnsupportedDeviceError:
112                _LOGGER.info("Skipping unsupported device %s %s", product.summary_info(), device.summary_info())
113                unsupported_devices_counter.increment(device.pv or "unknown")
114                continue
115            supported_devices_counter.increment(device.pv or "unknown")
116            start_tasks.append(new_device.start_connect())
117            new_devices[duid] = new_device
118
119        self._devices.update(new_devices)
120        await asyncio.gather(*start_tasks)
121        return list(self._devices.values())
122
123    async def get_device(self, duid: str) -> RoborockDevice | None:
124        """Get a specific device by DUID."""
125        return self._devices.get(duid)
126
127    async def get_devices(self) -> list[RoborockDevice]:
128        """Get all discovered devices."""
129        return list(self._devices.values())
130
131    async def close(self) -> None:
132        """Close all MQTT connections and clean up resources."""
133        tasks = [device.close() for device in self._devices.values()]
134        self._devices.clear()
135        tasks.append(self._mqtt_session.close())
136        await asyncio.gather(*tasks)
137
138    def diagnostic_data(self) -> Mapping[str, Any]:
139        """Return diagnostics information about the device manager."""
140        return {
141            "home_data": redact_device_data(self._home_data.as_dict()) if self._home_data else None,
142            "devices": [device.diagnostic_data() for device in self._devices.values()],
143            "diagnostics": self._diagnostics.as_dict(),
144        }
145
146
147@dataclass
148class UserParams:
149    """Parameters for creating a new session with Roborock devices.
150
151    These parameters include the username, user data for authentication,
152    and an optional base URL for the Roborock API. The `user_data` and `base_url`
153    parameters are obtained from `RoborockApiClient` during the login process.
154    """
155
156    username: str
157    """The username (email) used for logging in."""
158
159    user_data: UserData
160    """This is the user data containing authentication information."""
161
162    base_url: str | None = None
163    """Optional base URL for the Roborock API.
164
165    This is used to speed up connection times by avoiding the need to
166    discover the API base URL each time. If not provided, the API client
167    will attempt to discover it automatically which may take multiple requests.
168    """
169
170
171def create_web_api_wrapper(
172    user_params: UserParams,
173    *,
174    cache: Cache | None = None,
175    session: aiohttp.ClientSession | None = None,
176) -> UserWebApiClient:
177    """Create a home data API wrapper from an existing API client."""
178
179    # Note: This will auto discover the API base URL. This can be improved
180    # by caching this next to `UserData` if needed to avoid unnecessary API calls.
181    client = RoborockApiClient(username=user_params.username, base_url=user_params.base_url, session=session)
182
183    return UserWebApiClient(client, user_params.user_data)
184
185
186async def create_device_manager(
187    user_params: UserParams,
188    *,
189    cache: Cache | None = None,
190    map_parser_config: MapParserConfig | None = None,
191    session: aiohttp.ClientSession | None = None,
192    ready_callback: DeviceReadyCallback | None = None,
193    mqtt_session_unauthorized_hook: SessionUnauthorizedHook | None = None,
194    prefer_cache: bool = True,
195) -> DeviceManager:
196    """Convenience function to create and initialize a DeviceManager.
197
198    Args:
199        user_params: Parameters for creating the user session.
200        cache: Optional cache implementation to use for caching device data.
201        map_parser_config: Optional configuration for parsing maps.
202        session: Optional aiohttp ClientSession to use for HTTP requests.
203        ready_callback: Optional callback to be notified when a device is ready.
204        mqtt_session_unauthorized_hook: Optional hook for MQTT session unauthorized
205          events which may indicate rate limiting or revoked credentials. The
206          caller may use this to refresh authentication tokens as needed.
207        prefer_cache: Whether to prefer cached device data over always fetching it from the API.
208
209    Returns:
210        An initialized DeviceManager with discovered devices.
211    """
212    if cache is None:
213        cache = NoCache()
214
215    web_api = create_web_api_wrapper(user_params, session=session, cache=cache)
216    user_data = user_params.user_data
217
218    diagnostics = Diagnostics()
219
220    mqtt_params = create_mqtt_params(user_data.rriot)
221    mqtt_params.diagnostics = diagnostics.subkey("mqtt_session")
222    mqtt_params.unauthorized_hook = mqtt_session_unauthorized_hook
223    mqtt_session = await create_lazy_mqtt_session(mqtt_params)
224
225    def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
226        channel: Channel
227        trait: Trait
228        device_cache: DeviceCache = DeviceCache(device.duid, cache)
229        match device.pv:
230            case DeviceVersion.V1:
231                channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, device_cache)
232                trait = v1.create(
233                    device.duid,
234                    product,
235                    home_data,
236                    channel.rpc_channel,
237                    channel.mqtt_rpc_channel,
238                    channel.map_rpc_channel,
239                    web_api,
240                    device_cache=device_cache,
241                    map_parser_config=map_parser_config,
242                )
243            case DeviceVersion.A01:
244                channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
245                trait = a01.create(product, channel)
246            case DeviceVersion.B01:
247                channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
248                model_part = product.model.split(".")[-1]
249                if "ss" in model_part:
250                    trait = b01.q10.create(channel)
251                elif "sc" in model_part:
252                    # Q7 devices start with 'sc' in their model naming.
253                    trait = b01.q7.create(channel)
254                else:
255                    raise UnsupportedDeviceError(f"Device {device.name} has unsupported B01 model: {product.model}")
256            case _:
257                raise UnsupportedDeviceError(
258                    f"Device {device.name} has unsupported version {device.pv} {product.model}"
259                )
260
261        dev = RoborockDevice(device, product, channel, trait)
262        if ready_callback:
263            dev.add_ready_callback(ready_callback)
264        return dev
265
266    manager = DeviceManager(web_api, device_creator, mqtt_session=mqtt_session, cache=cache, diagnostics=diagnostics)
267    await manager.discover_devices(prefer_cache)
268    return manager