Overhaul Roborock integration to use new devices based API (home-assistant#154837)

Author: allenporter

homeassistant/components/roborock/__init__.py

@@ -9,16 +9,14 @@
 from typing import Any
 
 from roborock import (
-    HomeDataRoom,
     RoborockException,
     RoborockInvalidCredentials,
     RoborockInvalidUserAgreement,
     RoborockNoUserAgreement,
 )
-from roborock.data import DeviceData, HomeDataDevice, HomeDataProduct, UserData
-from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
-from roborock.version_a01_apis import RoborockMqttClientA01
-from roborock.web_api import RoborockApiClient
+from roborock.data import UserData
+from roborock.devices.device import RoborockDevice
+from roborock.devices.device_manager import UserParams, create_device_manager
 
 from homeassistant.const import CONF_USERNAME, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant
@@ -32,8 +30,10 @@
     RoborockCoordinators,
     RoborockDataUpdateCoordinator,
     RoborockDataUpdateCoordinatorA01,
+    RoborockWashingMachineUpdateCoordinator,
+    RoborockWetDryVacUpdateCoordinator,
 )
-from .roborock_storage import async_remove_map_storage
+from .roborock_storage import CacheStore, async_remove_map_storage
 
 SCAN_INTERVAL = timedelta(seconds=30)
 
@@ -44,14 +44,18 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
     """Set up roborock from a config entry."""
 
     user_data = UserData.from_dict(entry.data[CONF_USER_DATA])
-    api_client = RoborockApiClient(
-        entry.data[CONF_USERNAME],
-        entry.data[CONF_BASE_URL],
-        session=async_get_clientsession(hass),
+    user_params = UserParams(
+        username=entry.data[CONF_USERNAME],
+        user_data=user_data,
+        base_url=entry.data[CONF_BASE_URL],
     )
-    _LOGGER.debug("Getting home data")
+    cache = CacheStore(hass, entry.entry_id)
     try:
-        home_data = await api_client.get_home_data_v3(user_data)
+        device_manager = await create_device_manager(
+            user_params,
+            cache=cache,
+            session=async_get_clientsession(hass),
+        )
     except RoborockInvalidCredentials as err:
         raise ConfigEntryAuthFailed(
             "Invalid credentials",
@@ -75,29 +79,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
             translation_domain=DOMAIN,
             translation_key="home_data_fail",
         ) from err
+    devices = await device_manager.get_devices()
+    _LOGGER.debug("Device manager found %d devices", len(devices))
+    for device in devices:
+        entry.async_on_unload(device.close)
 
-    _LOGGER.debug("Got home data %s", home_data)
-    all_devices: list[HomeDataDevice] = home_data.devices + home_data.received_devices
-    device_map: dict[str, HomeDataDevice] = {
-        device.duid: device for device in all_devices
-    }
-    product_info: dict[str, HomeDataProduct] = {
-        product.id: product for product in home_data.products
-    }
-    # Get a Coordinator if the device is available or if we have connected to the device before
     coordinators = await asyncio.gather(
-        *build_setup_functions(
-            hass,
-            entry,
-            device_map,
-            user_data,
-            product_info,
-            home_data.rooms,
-            api_client,
-        ),
+        *build_setup_functions(hass, entry, devices, user_data),
         return_exceptions=True,
     )
-    # Valid coordinators are those where we had networking cached or we could get networking
     v1_coords = [
         coord
         for coord in coordinators
@@ -115,17 +105,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
             translation_key="no_coordinators",
         )
     valid_coordinators = RoborockCoordinators(v1_coords, a01_coords)
-    await asyncio.gather(
-        *(coord.refresh_coordinator_map() for coord in valid_coordinators.v1)
-    )
 
     async def on_stop(_: Any) -> None:
         _LOGGER.debug("Shutting down roborock")
         await asyncio.gather(
             *(
                 coordinator.async_shutdown()
                 for coordinator in valid_coordinators.values()
-            )
+            ),
+            cache.flush(),
         )
 
     entry.async_on_unload(
@@ -138,6 +126,17 @@ async def on_stop(_: Any) -> None:
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
+    _remove_stale_devices(hass, entry, devices)
+
+    return True
+
+
+def _remove_stale_devices(
+    hass: HomeAssistant,
+    entry: RoborockConfigEntry,
+    devices: list[RoborockDevice],
+) -> None:
+    device_map: dict[str, RoborockDevice] = {device.duid: device for device in devices}
     device_registry = dr.async_get(hass)
     device_entries = dr.async_entries_for_config_entry(
         device_registry, config_entry_id=entry.entry_id
@@ -159,8 +158,6 @@ async def on_stop(_: Any) -> None:
             remove_config_entry_id=entry.entry_id,
         )
 
-    return True
-
 
 async def async_migrate_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> bool:
     """Migrate old configuration entries to the new format."""
@@ -190,11 +187,8 @@ async def async_migrate_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -
 def build_setup_functions(
     hass: HomeAssistant,
     entry: RoborockConfigEntry,
-    device_map: dict[str, HomeDataDevice],
+    devices: list[RoborockDevice],
     user_data: UserData,
-    product_info: dict[str, HomeDataProduct],
-    home_data_rooms: list[HomeDataRoom],
-    api_client: RoborockApiClient,
 ) -> list[
     Coroutine[
         Any,
@@ -203,134 +197,45 @@ def build_setup_functions(
     ]
 ]:
     """Create a list of setup functions that can later be called asynchronously."""
-    return [
-        setup_device(
-            hass,
-            entry,
-            user_data,
-            device,
-            product_info[device.product_id],
-            home_data_rooms,
-            api_client,
-        )
-        for device in device_map.values()
-    ]
-
+    coordinators: list[
+        RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01
+    ] = []
+    for device in devices:
+        _LOGGER.debug("Creating device %s: %s", device.name, device)
+        if device.v1_properties is not None:
+            coordinators.append(
+                RoborockDataUpdateCoordinator(hass, entry, device, device.v1_properties)
+            )
+        elif device.dyad is not None:
+            coordinators.append(
+                RoborockWetDryVacUpdateCoordinator(hass, entry, device, device.dyad)
+            )
+        elif device.zeo is not None:
+            coordinators.append(
+                RoborockWashingMachineUpdateCoordinator(hass, entry, device, device.zeo)
+            )
+        else:
+            _LOGGER.warning(
+                "Not adding device %s because its protocol version %s or category %s is not supported",
+                device.duid,
+                device.device_info.pv,
+                device.product.category.name,
+            )
 
-async def setup_device(
-    hass: HomeAssistant,
-    entry: RoborockConfigEntry,
-    user_data: UserData,
-    device: HomeDataDevice,
-    product_info: HomeDataProduct,
-    home_data_rooms: list[HomeDataRoom],
-    api_client: RoborockApiClient,
-) -> RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01 | None:
-    """Set up a coordinator for a given device."""
-    if device.pv == "1.0":
-        return await setup_device_v1(
-            hass, entry, user_data, device, product_info, home_data_rooms, api_client
-        )
-    if device.pv == "A01":
-        return await setup_device_a01(hass, entry, user_data, device, product_info)
-    _LOGGER.warning(
-        "Not adding device %s because its protocol version %s or category %s is not supported",
-        device.duid,
-        device.pv,
-        product_info.category.name,
-    )
-    return None
+    return [setup_coordinator(coordinator) for coordinator in coordinators]
 
 
-async def setup_device_v1(
-    hass: HomeAssistant,
-    entry: RoborockConfigEntry,
-    user_data: UserData,
-    device: HomeDataDevice,
-    product_info: HomeDataProduct,
-    home_data_rooms: list[HomeDataRoom],
-    api_client: RoborockApiClient,
-) -> RoborockDataUpdateCoordinator | None:
-    """Set up a device Coordinator."""
-    mqtt_client = await hass.async_add_executor_job(
-        RoborockMqttClientV1, user_data, DeviceData(device, product_info.model)
-    )
-    try:
-        await mqtt_client.async_connect()
-        networking = await mqtt_client.get_networking()
-        if networking is None:
-            # If the api does not return an error but does return None for
-            # get_networking - then we need to go through cache checking.
-            raise RoborockException("Networking request returned None.")  # noqa: TRY301
-    except RoborockException as err:
-        _LOGGER.warning(
-            "Not setting up %s because we could not get the network information of the device. "
-            "Please confirm it is online and the Roborock servers can communicate with it",
-            device.name,
-        )
-        _LOGGER.debug(err)
-        await mqtt_client.async_release()
-        raise
-    coordinator = RoborockDataUpdateCoordinator(
-        hass,
-        entry,
-        device,
-        networking,
-        product_info,
-        mqtt_client,
-        home_data_rooms,
-        api_client,
-        user_data,
-    )
+async def setup_coordinator(
+    coordinator: RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01,
+) -> RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01 | None:
+    """Set up a single coordinator."""
     try:
         await coordinator.async_config_entry_first_refresh()
-    except ConfigEntryNotReady as ex:
+    except ConfigEntryNotReady:
         await coordinator.async_shutdown()
-        if isinstance(coordinator.api, RoborockMqttClientV1):
-            _LOGGER.warning(
-                "Not setting up %s because the we failed to get data for the first time using the online client. "
-                "Please ensure your Home Assistant instance can communicate with this device. "
-                "You may need to open firewall instances on your Home Assistant network and on your Vacuum's network",
-                device.name,
-            )
-            # Most of the time if we fail to connect using the mqtt client, the problem is due to firewall,
-            # but in case if it isn't, the error can be included in debug logs for the user to grab.
-            if coordinator.last_exception:
-                _LOGGER.debug(coordinator.last_exception)
-                raise coordinator.last_exception from ex
-        elif coordinator.last_exception:
-            # If this is reached, we have verified that we can communicate with the Vacuum locally,
-            # so if there is an error here - it is not a communication issue but some other problem
-            extra_error = f"Please create an issue with the following error included: {coordinator.last_exception}"
-            _LOGGER.warning(
-                "Not setting up %s because the coordinator failed to get data for the first time using the "
-                "offline client %s",
-                device.name,
-                extra_error,
-            )
-            raise coordinator.last_exception from ex
-    return coordinator
-
-
-async def setup_device_a01(
-    hass: HomeAssistant,
-    entry: RoborockConfigEntry,
-    user_data: UserData,
-    device: HomeDataDevice,
-    product_info: HomeDataProduct,
-) -> RoborockDataUpdateCoordinatorA01 | None:
-    """Set up a A01 protocol device."""
-    mqtt_client = await hass.async_add_executor_job(
-        RoborockMqttClientA01,
-        user_data,
-        DeviceData(device, product_info.model),
-        product_info.category,
-    )
-    coord = RoborockDataUpdateCoordinatorA01(
-        hass, entry, device, product_info, mqtt_client
-    )
-    await coord.async_config_entry_first_refresh()
-    return coord
+        raise
+    else:
+        return coordinator
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> bool:
@@ -341,3 +246,5 @@ async def async_unload_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
 async def async_remove_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> None:
     """Handle removal of an entry."""
     await async_remove_map_storage(hass, entry.entry_id)
+    store = CacheStore(hass, entry.entry_id)
+    await store.async_remove()

homeassistant/components/roborock/binary_sensor.py

@@ -6,7 +6,6 @@
 from dataclasses import dataclass
 
 from roborock.data import RoborockStateCode
-from roborock.roborock_typing import DeviceProp
 
 from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
@@ -19,6 +18,7 @@
 
 from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
 from .entity import RoborockCoordinatedEntityV1
+from .models import DeviceState
 
 PARALLEL_UPDATES = 0
 
@@ -27,9 +27,11 @@
 class RoborockBinarySensorDescription(BinarySensorEntityDescription):
     """A class that describes Roborock binary sensors."""
 
-    value_fn: Callable[[DeviceProp], bool | int | None]
-    # If it is a dock entity
+    value_fn: Callable[[DeviceState], bool | int | None]
+    """A function that extracts the sensor value from DeviceState."""
+
     is_dock_entity: bool = False
+    """Whether this sensor is for the dock."""
 
 
 BINARY_SENSOR_DESCRIPTIONS = [
@@ -92,7 +94,7 @@ async def async_setup_entry(
         )
         for coordinator in config_entry.runtime_data.v1
         for description in BINARY_SENSOR_DESCRIPTIONS
-        if description.value_fn(coordinator.roborock_device_info.props) is not None
+        if description.value_fn(coordinator.data) is not None
     )
 
 
@@ -117,8 +119,4 @@ def __init__(
     @property
     def is_on(self) -> bool:
         """Return the value reported by the sensor."""
-        return bool(
-            self.entity_description.value_fn(
-                self.coordinator.roborock_device_info.props
-            )
-        )
+        return bool(self.entity_description.value_fn(self.coordinator.data))