Explicitly pass in the config_entry in roborock coordinator (#137970)

explicitly pass in the config_entry in coordinator

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
 

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 RoborockConfigEntry, RoborockDataUpdateCoordinator
-from .coordinator import RoborockDataUpdateCoordinator
 from .entity import RoborockCoordinatedEntityV1
 
 

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 RoborockConfigEntry, RoborockDataUpdateCoordinator
-from .coordinator import RoborockDataUpdateCoordinator
 from .entity import RoborockEntityV1
 
 

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,

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"]
 

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
 
 

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 .const import DOMAIN
-from .coordinator import RoborockDataUpdateCoordinator
+from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
 from .entity import RoborockEntityV1
 
 _LOGGER = logging.getLogger(__name__)

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
 
 

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 (
-from .coordinator import RoborockDataUpdateCoordinator, RoborockDataUpdateCoordinatorA01
+    RoborockConfigEntry,
+    RoborockDataUpdateCoordinator,
+    RoborockDataUpdateCoordinatorA01,
+)
 from .entity import RoborockCoordinatedEntityA01, RoborockCoordinatedEntityV1
 
 

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 .const import DOMAIN
-from .coordinator import RoborockDataUpdateCoordinator
+from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
 from .entity import RoborockEntityV1
 
 _LOGGER = logging.getLogger(__name__)

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 .const import DOMAIN
-from .coordinator import RoborockDataUpdateCoordinator
+from .coordinator import RoborockConfigEntry, RoborockDataUpdateCoordinator
 from .entity import RoborockEntityV1
 
 _LOGGER = logging.getLogger(__name__)

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