diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index 310c5fee92b9b9b0393edd380d37ab6682f6600c..50ffbaaa6e1f0d492a7e595d7331103605dd053c 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 2fd6dd8d7d5edf3a00b2c8afdd0c58f79810d008..779d3ee234de4335cb40c1a2c1e84a930c739360 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 445033a0f6dada90e75a1b41172a4a8529159456..50d84e37a44883849b747abb28b2dd704e44a8aa 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 9be8b6f4d630a4e33ae1c1d7938010e2fea10d37..63de0da6a7f7ada392a5c5005e508ecb294a5501 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 d1731d289dbbc169e2000c2ba304b650be10f66b..4ead7e9635d6dad99f9d9c2451bd383792570169 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 5e776d40f2d32f4d822f63afbeed9ca7ba4b81a1..e86f07ad2048a5e20cbe2b0b66461bcfdea97fb1 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 c6073645086862ec6cbb05f349c3491fa0e953b7..8966652c24dd7b948965ac6e6180faada7821ef6 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 3be7461d1494d19238cdb90cdc520486b266c4e2..71c996f0b53f2bba7c3ac572e2ee23c7a56db617 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 cdfc0c2dc9664336ca42d1bf201ce56c1d3d69d0..6cc562fb533b0969c55670b343a1370d61f73cd1 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 21ab26c001387800e7a9112ac8a6297c2cdfd864..b0fbb18ed5606a976737c0dff332cc09d1a39dc8 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 cefcc85d7f81b82f430fc1815d5d52702c7f56ad..90f5002a23eccf60abe81a7ec000321e44fa3036 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 0437ce781f1a4b7bcf046f24e6d39cee73a97ca3..704f093d3fd1dd6aed3b7d75e74576e0ff072369 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(