From 31e9de3b95803762083a238a8469a328d768e9f9 Mon Sep 17 00:00:00 2001
From: Robert Resch <robert@resch.dev>
Date: Wed, 26 Jun 2024 19:42:15 +0200
Subject: [PATCH] Adapt Roborock to runtime_data (#120578)
* Adopt Roborock to runtime_data
* Fix name
---
homeassistant/components/roborock/__init__.py | 27 +++++++++++--------
.../components/roborock/binary_sensor.py | 9 +++----
homeassistant/components/roborock/button.py | 9 +++----
.../components/roborock/diagnostics.py | 8 +++---
homeassistant/components/roborock/image.py | 8 +++---
homeassistant/components/roborock/number.py | 9 +++----
homeassistant/components/roborock/select.py | 9 +++----
homeassistant/components/roborock/sensor.py | 8 +++---
homeassistant/components/roborock/switch.py | 9 +++----
homeassistant/components/roborock/time.py | 9 +++----
homeassistant/components/roborock/vacuum.py | 8 +++---
tests/components/roborock/test_init.py | 1 -
12 files changed, 46 insertions(+), 68 deletions(-)
diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py
index 310c5fee92b..50ffbaaa6e1 100644
--- a/homeassistant/components/roborock/__init__.py
+++ b/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):
- 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
+ return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-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)
diff --git a/homeassistant/components/roborock/binary_sensor.py b/homeassistant/components/roborock/binary_sensor.py
index 2fd6dd8d7d5..779d3ee234d 100644
--- a/homeassistant/components/roborock/binary_sensor.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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
)
diff --git a/homeassistant/components/roborock/button.py b/homeassistant/components/roborock/button.py
index 445033a0f6d..50d84e37a44 100644
--- a/homeassistant/components/roborock/button.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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)
)
diff --git a/homeassistant/components/roborock/diagnostics.py b/homeassistant/components/roborock/diagnostics.py
index 9be8b6f4d63..63de0da6a7f 100644
--- a/homeassistant/components/roborock/diagnostics.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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),
diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py
index d1731d289db..4ead7e9635d 100644
--- a/homeassistant/components/roborock/image.py
+++ b/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
)
)
)
diff --git a/homeassistant/components/roborock/number.py b/homeassistant/components/roborock/number.py
index 5e776d40f2d..e86f07ad204 100644
--- a/homeassistant/components/roborock/number.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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.
diff --git a/homeassistant/components/roborock/select.py b/homeassistant/components/roborock/select.py
index c6073645086..8966652c24d 100644
--- a/homeassistant/components/roborock/select.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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(
diff --git a/homeassistant/components/roborock/sensor.py b/homeassistant/components/roborock/sensor.py
index 3be7461d149..71c996f0b53 100644
--- a/homeassistant/components/roborock/sensor.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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,
diff --git a/homeassistant/components/roborock/switch.py b/homeassistant/components/roborock/switch.py
index cdfc0c2dc96..6cc562fb533 100644
--- a/homeassistant/components/roborock/switch.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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.
diff --git a/homeassistant/components/roborock/time.py b/homeassistant/components/roborock/time.py
index 21ab26c0013..b0fbb18ed56 100644
--- a/homeassistant/components/roborock/time.py
+++ b/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 .const import DOMAIN
+from . import RoborockConfigEntry
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.
diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py
index cefcc85d7f8..90f5002a23e 100644
--- a/homeassistant/components/roborock/vacuum.py
+++ b/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)
)
diff --git a/tests/components/roborock/test_init.py b/tests/components/roborock/test_init.py
index 0437ce781f1..704f093d3fd 100644
--- a/tests/components/roborock/test_init.py
+++ b/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(
--
GitLab