From a526baa8316e8bfce5c158749ef64640f556c30a Mon Sep 17 00:00:00 2001 From: Michael <35783820+mib1185@users.noreply.github.com> Date: Sun, 9 Feb 2025 05:53:44 +0100 Subject: [PATCH] Explicitly pass in the config_entry in roborock coordinator (#137970) explicitly pass in the config_entry in coordinator --- homeassistant/components/roborock/__init__.py | 48 +++++++++---------- .../components/roborock/binary_sensor.py | 3 +- homeassistant/components/roborock/button.py | 3 +- .../components/roborock/coordinator.py | 40 ++++++++++++++-- .../components/roborock/diagnostics.py | 2 +- homeassistant/components/roborock/image.py | 3 +- homeassistant/components/roborock/number.py | 4 +- homeassistant/components/roborock/select.py | 3 +- homeassistant/components/roborock/sensor.py | 7 ++- homeassistant/components/roborock/switch.py | 4 +- homeassistant/components/roborock/time.py | 4 +- homeassistant/components/roborock/vacuum.py | 3 +- 12 files changed, 77 insertions(+), 47 deletions(-) diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index b383c1acfd7..764518df636 100644 --- a/homeassistant/components/roborock/__init__.py +++ b/homeassistant/components/roborock/__init__.py @@ -4,7 +4,6 @@ from __future__ import annotations import asyncio from collections.abc import Coroutine -from dataclasses import dataclass from datetime import timedelta import logging from typing import Any @@ -21,35 +20,23 @@ 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 homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_USERNAME, EVENT_HOMEASSISTANT_STOP from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from .const import CONF_BASE_URL, CONF_USER_DATA, DOMAIN, PLATFORMS -from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01 +from .coordinator import ( + RoborockConfigEntry, + RoborockCoordinators, + RoborockDataUpdateCoordinator, + RoborockDataUpdateCoordinatorA01, +) from .roborock_storage import async_remove_map_storage SCAN_INTERVAL = timedelta(seconds=30) _LOGGER = logging.getLogger(__name__) -type RoborockConfigEntry = ConfigEntry[RoborockCoordinators] - - -@dataclass -class RoborockCoordinators: - """Roborock coordinators type.""" - - v1: list[RoborockDataUpdateCoordinator] - a01: list[RoborockDataUpdateCoordinatorA01] - - def values( - self, - ) -> list[RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01]: - """Return all coordinators.""" - return self.v1 + self.a01 - async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> bool: """Set up roborock from a config entry.""" @@ -95,7 +82,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> # 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, device_map, user_data, product_info, home_data.rooms + hass, entry, device_map, user_data, product_info, home_data.rooms ), return_exceptions=True, ) @@ -142,6 +129,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) -> def build_setup_functions( hass: HomeAssistant, + entry: RoborockConfigEntry, device_map: dict[str, HomeDataDevice], user_data: UserData, product_info: dict[str, HomeDataProduct], @@ -156,7 +144,12 @@ def build_setup_functions( """Create a list of setup functions that can later be called asynchronously.""" return [ setup_device( - hass, user_data, device, product_info[device.product_id], home_data_rooms + hass, + entry, + user_data, + device, + product_info[device.product_id], + home_data_rooms, ) for device in device_map.values() ] @@ -164,6 +157,7 @@ def build_setup_functions( async def setup_device( hass: HomeAssistant, + entry: RoborockConfigEntry, user_data: UserData, device: HomeDataDevice, product_info: HomeDataProduct, @@ -172,10 +166,10 @@ async def setup_device( """Set up a coordinator for a given device.""" if device.pv == "1.0": return await setup_device_v1( - hass, user_data, device, product_info, home_data_rooms + hass, entry, user_data, device, product_info, home_data_rooms ) if device.pv == "A01": - return await setup_device_a01(hass, user_data, device, product_info) + 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, @@ -187,6 +181,7 @@ async def setup_device( async def setup_device_v1( hass: HomeAssistant, + entry: RoborockConfigEntry, user_data: UserData, device: HomeDataDevice, product_info: HomeDataProduct, @@ -212,7 +207,7 @@ async def setup_device_v1( await mqtt_client.async_release() raise coordinator = RoborockDataUpdateCoordinator( - hass, device, networking, product_info, mqtt_client, home_data_rooms + hass, entry, device, networking, product_info, mqtt_client, home_data_rooms ) try: await coordinator.async_config_entry_first_refresh() @@ -246,6 +241,7 @@ async def setup_device_v1( async def setup_device_a01( hass: HomeAssistant, + entry: RoborockConfigEntry, user_data: UserData, device: HomeDataDevice, product_info: HomeDataProduct, @@ -254,7 +250,9 @@ async def setup_device_a01( mqtt_client = RoborockMqttClientA01( user_data, DeviceData(device, product_info.name), product_info.category ) - coord = RoborockDataUpdateCoordinatorA01(hass, device, product_info, mqtt_client) + coord = RoborockDataUpdateCoordinatorA01( + hass, entry, device, product_info, mqtt_client + ) await coord.async_config_entry_first_refresh() return coord diff --git a/homeassistant/components/roborock/binary_sensor.py b/homeassistant/components/roborock/binary_sensor.py index b88556ea857..c734eaf5ce8 100644 --- a/homeassistant/components/roborock/binary_sensor.py +++ b/homeassistant/components/roborock/binary_sensor.py @@ -16,8 +16,7 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RoborockConfigEntry -from .coordinator import RoborockDataUpdateCoordinator +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockCoordinatedEntityV1 diff --git a/homeassistant/components/roborock/button.py b/homeassistant/components/roborock/button.py index 2f214c7c51c..038f224f726 100644 --- a/homeassistant/components/roborock/button.py +++ b/homeassistant/components/roborock/button.py @@ -11,8 +11,7 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RoborockConfigEntry -from .coordinator import RoborockDataUpdateCoordinator +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockEntityV1 diff --git a/homeassistant/components/roborock/coordinator.py b/homeassistant/components/roborock/coordinator.py index 8860a5c1f43..918c7159ee3 100644 --- a/homeassistant/components/roborock/coordinator.py +++ b/homeassistant/components/roborock/coordinator.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from dataclasses import dataclass from datetime import timedelta import logging @@ -35,14 +36,32 @@ SCAN_INTERVAL = timedelta(seconds=30) _LOGGER = logging.getLogger(__name__) +@dataclass +class RoborockCoordinators: + """Roborock coordinators type.""" + + v1: list[RoborockDataUpdateCoordinator] + a01: list[RoborockDataUpdateCoordinatorA01] + + def values( + self, + ) -> list[RoborockDataUpdateCoordinator | RoborockDataUpdateCoordinatorA01]: + """Return all coordinators.""" + return self.v1 + self.a01 + + +type RoborockConfigEntry = ConfigEntry[RoborockCoordinators] + + class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): """Class to manage fetching data from the API.""" - config_entry: ConfigEntry + config_entry: RoborockConfigEntry def __init__( self, hass: HomeAssistant, + config_entry: RoborockConfigEntry, device: HomeDataDevice, device_networking: NetworkInfo, product_info: HomeDataProduct, @@ -50,7 +69,13 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]): home_data_rooms: list[HomeDataRoom], ) -> None: """Initialize.""" - super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL) + super().__init__( + hass, + _LOGGER, + config_entry=config_entry, + name=DOMAIN, + update_interval=SCAN_INTERVAL, + ) self.roborock_device_info = RoborockHassDeviceInfo( device, device_networking, @@ -186,15 +211,24 @@ class RoborockDataUpdateCoordinatorA01( ): """Class to manage fetching data from the API for A01 devices.""" + config_entry: RoborockConfigEntry + def __init__( self, hass: HomeAssistant, + config_entry: RoborockConfigEntry, device: HomeDataDevice, product_info: HomeDataProduct, api: RoborockClientA01, ) -> None: """Initialize.""" - super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL) + super().__init__( + hass, + _LOGGER, + config_entry=config_entry, + name=DOMAIN, + update_interval=SCAN_INTERVAL, + ) self.api = api self.device_info = DeviceInfo( name=device.name, diff --git a/homeassistant/components/roborock/diagnostics.py b/homeassistant/components/roborock/diagnostics.py index e784e4ce837..4602b4bd02a 100644 --- a/homeassistant/components/roborock/diagnostics.py +++ b/homeassistant/components/roborock/diagnostics.py @@ -8,7 +8,7 @@ from homeassistant.components.diagnostics import async_redact_data from homeassistant.const import CONF_UNIQUE_ID from homeassistant.core import HomeAssistant -from . import RoborockConfigEntry +from .coordinator import RoborockConfigEntry TO_REDACT_CONFIG = ["token", "sn", "rruid", CONF_UNIQUE_ID, "username", "uid"] diff --git a/homeassistant/components/roborock/image.py b/homeassistant/components/roborock/image.py index b4776c27164..ff1c94957e0 100644 --- a/homeassistant/components/roborock/image.py +++ b/homeassistant/components/roborock/image.py @@ -19,7 +19,6 @@ from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util import dt as dt_util -from . import RoborockConfigEntry from .const import ( DEFAULT_DRAWABLES, DOMAIN, @@ -28,7 +27,7 @@ from .const import ( MAP_FILE_FORMAT, MAP_SLEEP, ) -from .coordinator import RoborockDataUpdateCoordinator +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockCoordinatedEntityV1 diff --git a/homeassistant/components/roborock/number.py b/homeassistant/components/roborock/number.py index 7f568ae824b..97aa8c2ffd4 100644 --- a/homeassistant/components/roborock/number.py +++ b/homeassistant/components/roborock/number.py @@ -16,8 +16,8 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import DOMAIN, RoborockConfigEntry -from .coordinator import RoborockDataUpdateCoordinator +from .const import DOMAIN +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockEntityV1 _LOGGER = logging.getLogger(__name__) diff --git a/homeassistant/components/roborock/select.py b/homeassistant/components/roborock/select.py index 73cb95d2d7c..826af3e24e8 100644 --- a/homeassistant/components/roborock/select.py +++ b/homeassistant/components/roborock/select.py @@ -13,9 +13,8 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RoborockConfigEntry from .const import MAP_SLEEP -from .coordinator import RoborockDataUpdateCoordinator +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockCoordinatedEntityV1 diff --git a/homeassistant/components/roborock/sensor.py b/homeassistant/components/roborock/sensor.py index e01a03d7720..0d376debcbf 100644 --- a/homeassistant/components/roborock/sensor.py +++ b/homeassistant/components/roborock/sensor.py @@ -31,8 +31,11 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType -from . import RoborockConfigEntry -from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01 +from .coordinator import ( + RoborockConfigEntry, + RoborockDataUpdateCoordinator, + RoborockDataUpdateCoordinatorA01, +) from .entity import RoborockCoordinatedEntityA01, RoborockCoordinatedEntityV1 diff --git a/homeassistant/components/roborock/switch.py b/homeassistant/components/roborock/switch.py index b0c8c880188..ebf8225b4f5 100644 --- a/homeassistant/components/roborock/switch.py +++ b/homeassistant/components/roborock/switch.py @@ -18,8 +18,8 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import DOMAIN, RoborockConfigEntry -from .coordinator import RoborockDataUpdateCoordinator +from .const import DOMAIN +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockEntityV1 _LOGGER = logging.getLogger(__name__) diff --git a/homeassistant/components/roborock/time.py b/homeassistant/components/roborock/time.py index 1dd681dff1f..76f20bc6607 100644 --- a/homeassistant/components/roborock/time.py +++ b/homeassistant/components/roborock/time.py @@ -18,8 +18,8 @@ from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import DOMAIN, RoborockConfigEntry -from .coordinator import RoborockDataUpdateCoordinator +from .const import DOMAIN +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockEntityV1 _LOGGER = logging.getLogger(__name__) diff --git a/homeassistant/components/roborock/vacuum.py b/homeassistant/components/roborock/vacuum.py index 7582dadad16..e604ab6a209 100644 --- a/homeassistant/components/roborock/vacuum.py +++ b/homeassistant/components/roborock/vacuum.py @@ -18,14 +18,13 @@ from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import config_validation as cv, entity_platform from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RoborockConfigEntry from .const import ( DOMAIN, GET_MAPS_SERVICE_NAME, GET_VACUUM_CURRENT_POSITION_SERVICE_NAME, SET_VACUUM_GOTO_POSITION_SERVICE_NAME, ) -from .coordinator import RoborockDataUpdateCoordinator +from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator from .entity import RoborockCoordinatedEntityV1 from .image import ColorsPalette, ImageConfig, RoborockMapDataParser, Sizes