From c7c234c5dd6276675d744b276b4e432235165790 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Tue, 28 Jan 2025 11:34:57 +0100 Subject: [PATCH] Use runtime_data in electric_kiwi (#136699) --- .../components/electric_kiwi/__init__.py | 28 ++++++++--------- .../components/electric_kiwi/const.py | 3 -- .../components/electric_kiwi/coordinator.py | 31 +++++++++++++++++-- .../components/electric_kiwi/select.py | 13 ++++---- .../components/electric_kiwi/sensor.py | 16 +++++----- 5 files changed, 56 insertions(+), 35 deletions(-) diff --git a/homeassistant/components/electric_kiwi/__init__.py b/homeassistant/components/electric_kiwi/__init__.py index 8c9a0b3950e..de8d87553a3 100644 --- a/homeassistant/components/electric_kiwi/__init__.py +++ b/homeassistant/components/electric_kiwi/__init__.py @@ -6,23 +6,25 @@ import aiohttp from electrickiwi_api import ElectricKiwiApi from electrickiwi_api.exceptions import ApiException -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow from . import api -from .const import ACCOUNT_COORDINATOR, DOMAIN, HOP_COORDINATOR from .coordinator import ( ElectricKiwiAccountDataCoordinator, + ElectricKiwiConfigEntry, ElectricKiwiHOPDataCoordinator, + ElectricKiwiRuntimeData, ) PLATFORMS: list[Platform] = [Platform.SELECT, Platform.SENSOR] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry( + hass: HomeAssistant, entry: ElectricKiwiConfigEntry +) -> bool: """Set up Electric Kiwi from a config entry.""" implementation = ( await config_entry_oauth2_flow.async_get_config_entry_implementation( @@ -44,8 +46,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ek_api = ElectricKiwiApi( api.AsyncConfigEntryAuth(aiohttp_client.async_get_clientsession(hass), session) ) - hop_coordinator = ElectricKiwiHOPDataCoordinator(hass, ek_api) - account_coordinator = ElectricKiwiAccountDataCoordinator(hass, ek_api) + hop_coordinator = ElectricKiwiHOPDataCoordinator(hass, entry, ek_api) + account_coordinator = ElectricKiwiAccountDataCoordinator(hass, entry, ek_api) try: await ek_api.set_active_session() @@ -54,19 +56,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: except ApiException as err: raise ConfigEntryNotReady from err - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - HOP_COORDINATOR: hop_coordinator, - ACCOUNT_COORDINATOR: account_coordinator, - } + entry.runtime_data = ElectricKiwiRuntimeData( + hop=hop_coordinator, account=account_coordinator + ) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry( + hass: HomeAssistant, entry: ElectricKiwiConfigEntry +) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/electric_kiwi/const.py b/homeassistant/components/electric_kiwi/const.py index 0b455b045cf..907b6247172 100644 --- a/homeassistant/components/electric_kiwi/const.py +++ b/homeassistant/components/electric_kiwi/const.py @@ -9,6 +9,3 @@ OAUTH2_TOKEN = "https://welcome.electrickiwi.co.nz/oauth/token" API_BASE_URL = "https://api.electrickiwi.co.nz" SCOPE_VALUES = "read_connection_detail read_billing_frequency read_account_running_balance read_consumption_summary read_consumption_averages read_hop_intervals_config read_hop_connection save_hop_connection read_session" - -HOP_COORDINATOR = "hop_coordinator" -ACCOUNT_COORDINATOR = "account_coordinator" diff --git a/homeassistant/components/electric_kiwi/coordinator.py b/homeassistant/components/electric_kiwi/coordinator.py index a10be5eafdd..2065da5d668 100644 --- a/homeassistant/components/electric_kiwi/coordinator.py +++ b/homeassistant/components/electric_kiwi/coordinator.py @@ -1,7 +1,10 @@ """Electric Kiwi coordinators.""" +from __future__ import annotations + import asyncio from collections import OrderedDict +from dataclasses import dataclass from datetime import timedelta import logging @@ -9,6 +12,7 @@ from electrickiwi_api import ElectricKiwiApi from electrickiwi_api.exceptions import ApiException, AuthException from electrickiwi_api.model import AccountBalance, Hop, HopIntervals +from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed @@ -19,14 +23,31 @@ ACCOUNT_SCAN_INTERVAL = timedelta(hours=6) HOP_SCAN_INTERVAL = timedelta(minutes=20) +@dataclass +class ElectricKiwiRuntimeData: + """ElectricKiwi runtime data.""" + + hop: ElectricKiwiHOPDataCoordinator + account: ElectricKiwiAccountDataCoordinator + + +type ElectricKiwiConfigEntry = ConfigEntry[ElectricKiwiRuntimeData] + + class ElectricKiwiAccountDataCoordinator(DataUpdateCoordinator[AccountBalance]): """ElectricKiwi Account Data object.""" - def __init__(self, hass: HomeAssistant, ek_api: ElectricKiwiApi) -> None: + def __init__( + self, + hass: HomeAssistant, + entry: ElectricKiwiConfigEntry, + ek_api: ElectricKiwiApi, + ) -> None: """Initialize ElectricKiwiAccountDataCoordinator.""" super().__init__( hass, _LOGGER, + config_entry=entry, name="Electric Kiwi Account Data", update_interval=ACCOUNT_SCAN_INTERVAL, ) @@ -48,11 +69,17 @@ class ElectricKiwiAccountDataCoordinator(DataUpdateCoordinator[AccountBalance]): class ElectricKiwiHOPDataCoordinator(DataUpdateCoordinator[Hop]): """ElectricKiwi HOP Data object.""" - def __init__(self, hass: HomeAssistant, ek_api: ElectricKiwiApi) -> None: + def __init__( + self, + hass: HomeAssistant, + entry: ElectricKiwiConfigEntry, + ek_api: ElectricKiwiApi, + ) -> None: """Initialize ElectricKiwiAccountDataCoordinator.""" super().__init__( hass, _LOGGER, + config_entry=entry, # Name of the data. For logging purposes. name="Electric Kiwi HOP Data", # Polling interval. Will only be polled if there are subscribers. diff --git a/homeassistant/components/electric_kiwi/select.py b/homeassistant/components/electric_kiwi/select.py index a3f073b8ca2..fa111381612 100644 --- a/homeassistant/components/electric_kiwi/select.py +++ b/homeassistant/components/electric_kiwi/select.py @@ -5,14 +5,13 @@ from __future__ import annotations import logging 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.helpers.update_coordinator import CoordinatorEntity -from .const import ATTRIBUTION, DOMAIN, HOP_COORDINATOR -from .coordinator import ElectricKiwiHOPDataCoordinator +from .const import ATTRIBUTION +from .coordinator import ElectricKiwiConfigEntry, ElectricKiwiHOPDataCoordinator _LOGGER = logging.getLogger(__name__) ATTR_EK_HOP_SELECT = "hop_select" @@ -25,12 +24,12 @@ HOP_SELECT = SelectEntityDescription( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: ElectricKiwiConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Electric Kiwi select setup.""" - hop_coordinator: ElectricKiwiHOPDataCoordinator = hass.data[DOMAIN][entry.entry_id][ - HOP_COORDINATOR - ] + hop_coordinator = entry.runtime_data.hop _LOGGER.debug("Setting up select entity") async_add_entities([ElectricKiwiSelectHOPEntity(hop_coordinator, HOP_SELECT)]) diff --git a/homeassistant/components/electric_kiwi/sensor.py b/homeassistant/components/electric_kiwi/sensor.py index 7672466106b..e070f9495c1 100644 --- a/homeassistant/components/electric_kiwi/sensor.py +++ b/homeassistant/components/electric_kiwi/sensor.py @@ -14,16 +14,16 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CURRENCY_DOLLAR, PERCENTAGE from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import dt as dt_util -from .const import ACCOUNT_COORDINATOR, ATTRIBUTION, DOMAIN, HOP_COORDINATOR +from .const import ATTRIBUTION from .coordinator import ( ElectricKiwiAccountDataCoordinator, + ElectricKiwiConfigEntry, ElectricKiwiHOPDataCoordinator, ) @@ -122,12 +122,12 @@ HOP_SENSOR_TYPES: tuple[ElectricKiwiHOPSensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: ElectricKiwiConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Electric Kiwi Sensors Setup.""" - account_coordinator: ElectricKiwiAccountDataCoordinator = hass.data[DOMAIN][ - entry.entry_id - ][ACCOUNT_COORDINATOR] + account_coordinator = entry.runtime_data.account entities: list[SensorEntity] = [ ElectricKiwiAccountEntity( @@ -137,9 +137,7 @@ async def async_setup_entry( for description in ACCOUNT_SENSOR_TYPES ] - hop_coordinator: ElectricKiwiHOPDataCoordinator = hass.data[DOMAIN][entry.entry_id][ - HOP_COORDINATOR - ] + hop_coordinator = entry.runtime_data.hop entities.extend( [ ElectricKiwiHOPEntity(hop_coordinator, description)