Explicitly pass in the config_entry in aosmith coordinator init (#137710)

explicitly pass in the config_entry in aosmith coordinator init

homeassistant/components/aosmith/__init__.py

@@ -2,31 +2,22 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-
 from py_aosmith import AOSmithAPIClient
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import aiohttp_client, device_registry as dr
 
 from .const import DOMAIN
-from .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator
+from .coordinator import (
+    AOSmithConfigEntry,
+    AOSmithData,
+    AOSmithEnergyCoordinator,
+    AOSmithStatusCoordinator,
+)
 
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.WATER_HEATER]
 
-type AOSmithConfigEntry = ConfigEntry[AOSmithData]
-
-
-@dataclass
-class AOSmithData:
-    """Data for the A. O. Smith integration."""
-
-    client: AOSmithAPIClient
-    status_coordinator: AOSmithStatusCoordinator
-    energy_coordinator: AOSmithEnergyCoordinator
-
 
 async def async_setup_entry(hass: HomeAssistant, entry: AOSmithConfigEntry) -> bool:
     """Set up A. O. Smith from a config entry."""
@@ -36,7 +27,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: AOSmithConfigEntry) -> b
     session = aiohttp_client.async_get_clientsession(hass)
     client = AOSmithAPIClient(email, password, session)
 
-    status_coordinator = AOSmithStatusCoordinator(hass, client)
+    status_coordinator = AOSmithStatusCoordinator(hass, entry, client)
     await status_coordinator.async_config_entry_first_refresh()
 
     device_registry = dr.async_get(hass)
@@ -53,7 +44,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: AOSmithConfigEntry) -> b
         )
 
     energy_coordinator = AOSmithEnergyCoordinator(
-        hass, client, list(status_coordinator.data)
+        hass, entry, client, list(status_coordinator.data)
     )
     await energy_coordinator.async_config_entry_first_refresh()
 

homeassistant/components/aosmith/coordinator.py

@@ -1,5 +1,6 @@
 """The data update coordinator for the A. O. Smith integration."""
 
+from dataclasses import dataclass
 import logging
 
 from py_aosmith import (
@@ -9,6 +10,7 @@ from py_aosmith import (
 )
 from py_aosmith.models import Device as AOSmithDevice
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
@@ -17,13 +19,37 @@ from .const import DOMAIN, ENERGY_USAGE_INTERVAL, FAST_INTERVAL, REGULAR_INTERVA
 
 _LOGGER = logging.getLogger(__name__)
 
+type AOSmithConfigEntry = ConfigEntry[AOSmithData]
+
+
+@dataclass
+class AOSmithData:
+    """Data for the A. O. Smith integration."""
+
+    client: AOSmithAPIClient
+    status_coordinator: "AOSmithStatusCoordinator"
+    energy_coordinator: "AOSmithEnergyCoordinator"
+
 
 class AOSmithStatusCoordinator(DataUpdateCoordinator[dict[str, AOSmithDevice]]):
     """Coordinator for device status, updating with a frequent interval."""
 
-    def __init__(self, hass: HomeAssistant, client: AOSmithAPIClient) -> None:
+    config_entry: AOSmithConfigEntry
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        config_entry: AOSmithConfigEntry,
+        client: AOSmithAPIClient,
+    ) -> None:
         """Initialize the coordinator."""
-        super().__init__(hass, _LOGGER, name=DOMAIN, update_interval=REGULAR_INTERVAL)
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=config_entry,
+            name=DOMAIN,
+            update_interval=REGULAR_INTERVAL,
+        )
         self.client = client
 
     async def _async_update_data(self) -> dict[str, AOSmithDevice]:
@@ -51,15 +77,22 @@ class AOSmithStatusCoordinator(DataUpdateCoordinator[dict[str, AOSmithDevice]]):
 class AOSmithEnergyCoordinator(DataUpdateCoordinator[dict[str, float]]):
     """Coordinator for energy usage data, updating with a slower interval."""
 
+    config_entry: AOSmithConfigEntry
+
     def __init__(
         self,
         hass: HomeAssistant,
+        config_entry: AOSmithConfigEntry,
         client: AOSmithAPIClient,
         junction_ids: list[str],
     ) -> None:
         """Initialize the coordinator."""
         super().__init__(
-            hass, _LOGGER, name=DOMAIN, update_interval=ENERGY_USAGE_INTERVAL
+            hass,
+            _LOGGER,
+            config_entry=config_entry,
+            name=DOMAIN,
+            update_interval=ENERGY_USAGE_INTERVAL,
         )
         self.client = client
         self.junction_ids = junction_ids

homeassistant/components/aosmith/diagnostics.py

@@ -7,7 +7,7 @@ from typing import Any
 from homeassistant.components.diagnostics import async_redact_data
 from homeassistant.core import HomeAssistant
 
-from . import AOSmithConfigEntry
+from .coordinator import AOSmithConfigEntry
 
 TO_REDACT = {
     "address",

homeassistant/components/aosmith/sensor.py

@@ -15,8 +15,11 @@ from homeassistant.const import PERCENTAGE, UnitOfEnergy
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import AOSmithConfigEntry
-from .coordinator import AOSmithEnergyCoordinator, AOSmithStatusCoordinator
+from .coordinator import (
+    AOSmithConfigEntry,
+    AOSmithEnergyCoordinator,
+    AOSmithStatusCoordinator,
+)
 from .entity import AOSmithEnergyEntity, AOSmithStatusEntity
 
 

homeassistant/components/aosmith/water_heater.py

@@ -17,8 +17,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import AOSmithConfigEntry
-from .coordinator import AOSmithStatusCoordinator
+from .coordinator import AOSmithConfigEntry, AOSmithStatusCoordinator
 from .entity import AOSmithStatusEntity
 
 MODE_HA_TO_AOSMITH = {