diff --git a/homeassistant/components/aosmith/__init__.py b/homeassistant/components/aosmith/__init__.py index dd60f69c4b9..7593365c573 100644 --- a/homeassistant/components/aosmith/__init__.py +++ b/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() diff --git a/homeassistant/components/aosmith/coordinator.py b/homeassistant/components/aosmith/coordinator.py index 3bf97e49cae..26029fee750 100644 --- a/homeassistant/components/aosmith/coordinator.py +++ b/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 diff --git a/homeassistant/components/aosmith/diagnostics.py b/homeassistant/components/aosmith/diagnostics.py index 94726731f75..4019bee4dc8 100644 --- a/homeassistant/components/aosmith/diagnostics.py +++ b/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", diff --git a/homeassistant/components/aosmith/sensor.py b/homeassistant/components/aosmith/sensor.py index b1c9852f647..8a7a98115fa 100644 --- a/homeassistant/components/aosmith/sensor.py +++ b/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 diff --git a/homeassistant/components/aosmith/water_heater.py b/homeassistant/components/aosmith/water_heater.py index f3dc8b3413f..110f997065b 100644 --- a/homeassistant/components/aosmith/water_heater.py +++ b/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 = {