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