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