Adapt Roborock to runtime_data (#120578)

* Adopt Roborock to runtime_data

* Fix name

homeassistant/components/roborock/__init__.py

@@ -28,6 +28,8 @@ SCAN_INTERVAL = timedelta(seconds=30)
 
 _LOGGER = logging.getLogger(__name__)
 
+type RoborockConfigEntry = ConfigEntry[RoborockCoordinators]
+
 
 @dataclass
 class RoborockCoordinators:
@@ -43,7 +45,7 @@ class RoborockCoordinators:
         return self.v1 + self.a01
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> bool:
     """Set up roborock from a config entry."""
 
     _LOGGER.debug("Integration async setup entry: %s", entry.as_dict())
@@ -99,7 +101,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             translation_key="no_coordinators",
         )
     valid_coordinators = RoborockCoordinators(v1_coords, a01_coords)
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = valid_coordinators
+
+    async def on_unload() -> None:
+        release_tasks = set()
+        for coordinator in valid_coordinators.values():
+            release_tasks.add(coordinator.release())
+        await asyncio.gather(*release_tasks)
+
+    entry.async_on_unload(on_unload)
+    entry.runtime_data = valid_coordinators
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
@@ -231,18 +242,12 @@ async def setup_device_a01(
     return coord
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> bool:
     """Handle removal of an entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-        release_tasks = set()
-        for coordinator in hass.data[DOMAIN][entry.entry_id].values():
-            release_tasks.add(coordinator.release())
-        hass.data[DOMAIN].pop(entry.entry_id)
-        await asyncio.gather(*release_tasks)
-    return unload_ok
 
 
-async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def update_listener(hass: HomeAssistant, entry: RoborockConfigEntry) -> None:
     """Handle options update."""
     # Reload entry to update data
     await hass.config_entries.async_reload(entry.entry_id)

homeassistant/components/roborock/binary_sensor.py

@@ -12,14 +12,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockCoordinatedEntityV1
 
@@ -72,17 +70,16 @@ BINARY_SENSOR_DESCRIPTIONS = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Roborock vacuum binary sensors."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     async_add_entities(
         RoborockBinarySensorEntity(
             coordinator,
             description,
         )
-        for coordinator in coordinators.v1
+        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
     )

homeassistant/components/roborock/button.py

@@ -7,14 +7,12 @@ from dataclasses import dataclass
 from roborock.roborock_typing import RoborockCommand
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockEntityV1
 
@@ -65,17 +63,16 @@ CONSUMABLE_BUTTON_DESCRIPTIONS = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roborock button platform."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     async_add_entities(
         RoborockButtonEntity(
             coordinator,
             description,
         )
-        for coordinator in coordinators.v1
+        for coordinator in config_entry.runtime_data.v1
         for description in CONSUMABLE_BUTTON_DESCRIPTIONS
         if isinstance(coordinator, RoborockDataUpdateCoordinator)
     )

homeassistant/components/roborock/diagnostics.py

@@ -5,12 +5,10 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics.util import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_UNIQUE_ID
 from homeassistant.core import HomeAssistant
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 
 TO_REDACT_CONFIG = ["token", "sn", "rruid", CONF_UNIQUE_ID, "username", "uid"]
 
@@ -18,10 +16,10 @@ TO_REDACT_COORD = ["duid", "localKey", "mac", "bssid"]
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: RoborockConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
+    coordinators = config_entry.runtime_data
 
     return {
         "config_entry": async_redact_data(config_entry.data, TO_REDACT_CONFIG),

homeassistant/components/roborock/image.py

@@ -13,7 +13,6 @@ from vacuum_map_parser_base.config.size import Sizes
 from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser
 
 from homeassistant.components.image import ImageEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
@@ -21,7 +20,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 import homeassistant.util.dt as dt_util
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
 from .const import DEFAULT_DRAWABLES, DOMAIN, DRAWABLES, IMAGE_CACHE_INTERVAL, MAP_SLEEP
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockCoordinatedEntityV1
@@ -29,12 +28,11 @@ from .device import RoborockCoordinatedEntityV1
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roborock image platform."""
 
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     drawables = [
         drawable
         for drawable, default_value in DEFAULT_DRAWABLES.items()
@@ -45,7 +43,7 @@ async def async_setup_entry(
             await asyncio.gather(
                 *(
                     create_coordinator_maps(coord, drawables)
-                    for coord in coordinators.v1
+                    for coord in config_entry.runtime_data.v1
                 )
             )
         )

homeassistant/components/roborock/number.py

@@ -11,14 +11,12 @@ from roborock.exceptions import RoborockException
 from roborock.version_1_apis.roborock_client_v1 import AttributeCache
 
 from homeassistant.components.number import NumberEntity, NumberEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockEntityV1
 
@@ -51,16 +49,15 @@ NUMBER_DESCRIPTIONS: list[RoborockNumberDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roborock number platform."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     possible_entities: list[
         tuple[RoborockDataUpdateCoordinator, RoborockNumberDescription]
     ] = [
         (coordinator, description)
-        for coordinator in coordinators.v1
+        for coordinator in config_entry.runtime_data.v1
         for description in NUMBER_DESCRIPTIONS
     ]
     # We need to check if this function is supported by the device.

homeassistant/components/roborock/select.py

@@ -8,14 +8,12 @@ from roborock.roborock_message import RoborockDataProtocol
 from roborock.roborock_typing import RoborockCommand
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockCoordinatedEntityV1
 
@@ -65,15 +63,14 @@ SELECT_DESCRIPTIONS: list[RoborockSelectDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roborock select platform."""
 
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     async_add_entities(
         RoborockSelectEntity(coordinator, description, options)
-        for coordinator in coordinators.v1
+        for coordinator in config_entry.runtime_data.v1
         for description in SELECT_DESCRIPTIONS
         if (
             options := description.options_lambda(

homeassistant/components/roborock/sensor.py

@@ -21,7 +21,6 @@ from homeassistant.components.sensor import (
     SensorEntity,
     SensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     AREA_SQUARE_METERS,
     PERCENTAGE,
@@ -33,8 +32,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01
 from .device import RoborockCoordinatedEntityA01, RoborockCoordinatedEntityV1
 
@@ -255,11 +253,11 @@ A01_SENSOR_DESCRIPTIONS: list[RoborockSensorDescriptionA01] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Roborock vacuum sensors."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
+    coordinators = config_entry.runtime_data
     async_add_entities(
         RoborockSensorEntity(
             coordinator,

homeassistant/components/roborock/switch.py

@@ -12,14 +12,12 @@ from roborock.command_cache import CacheableAttribute
 from roborock.version_1_apis.roborock_client_v1 import AttributeCache
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockEntityV1
 
@@ -99,16 +97,15 @@ SWITCH_DESCRIPTIONS: list[RoborockSwitchDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roborock switch platform."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     possible_entities: list[
         tuple[RoborockDataUpdateCoordinator, RoborockSwitchDescription]
     ] = [
         (coordinator, description)
-        for coordinator in coordinators.v1
+        for coordinator in config_entry.runtime_data.v1
         for description in SWITCH_DESCRIPTIONS
     ]
     # We need to check if this function is supported by the device.

homeassistant/components/roborock/time.py

@@ -13,14 +13,12 @@ from roborock.exceptions import RoborockException
 from roborock.version_1_apis.roborock_client_v1 import AttributeCache
 
 from homeassistant.components.time import TimeEntity, TimeEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
-from .const import DOMAIN
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockEntityV1
 
@@ -115,16 +113,15 @@ TIME_DESCRIPTIONS: list[RoborockTimeDescription] = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Roborock time platform."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     possible_entities: list[
         tuple[RoborockDataUpdateCoordinator, RoborockTimeDescription]
     ] = [
         (coordinator, description)
-        for coordinator in coordinators.v1
+        for coordinator in config_entry.runtime_data.v1
         for description in TIME_DESCRIPTIONS
     ]
     # We need to check if this function is supported by the device.

homeassistant/components/roborock/vacuum.py

@@ -17,13 +17,12 @@ from homeassistant.components.vacuum import (
     StateVacuumEntity,
     VacuumEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, ServiceResponse, SupportsResponse
 from homeassistant.helpers import entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
-from . import RoborockCoordinators
+from . import RoborockConfigEntry
 from .const import DOMAIN, GET_MAPS_SERVICE_NAME
 from .coordinator import RoborockDataUpdateCoordinator
 from .device import RoborockCoordinatedEntityV1
@@ -57,14 +56,13 @@ STATE_CODE_TO_STATE = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: RoborockConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Roborock sensor."""
-    coordinators: RoborockCoordinators = hass.data[DOMAIN][config_entry.entry_id]
     async_add_entities(
         RoborockVacuum(coordinator)
-        for coordinator in coordinators.v1
+        for coordinator in config_entry.runtime_data.v1
         if isinstance(coordinator, RoborockDataUpdateCoordinator)
     )
 

tests/components/roborock/test_init.py

@@ -29,7 +29,6 @@ async def test_unload_entry(
         await hass.async_block_till_done()
         assert mock_disconnect.call_count == 2
         assert setup_entry.state is ConfigEntryState.NOT_LOADED
-        assert not hass.data.get(DOMAIN)
 
 
 async def test_config_entry_not_ready(