From 31e9de3b95803762083a238a8469a328d768e9f9 Mon Sep 17 00:00:00 2001 From: Robert Resch 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(