Use runtime_data in electric_kiwi (#136699)

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)
+    hop_coordinator = ElectricKiwiHOPDataCoordinator(hass, entry, ek_api)
-    account_coordinator = ElectricKiwiAccountDataCoordinator(hass, 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] = {
+    entry.runtime_data = ElectricKiwiRuntimeData(
-        HOP_COORDINATOR: hop_coordinator,
+        hop=hop_coordinator, account=account_coordinator
-        ACCOUNT_COORDINATOR: 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):
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok

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"

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.

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 .const import ATTRIBUTION
-from .coordinator import ElectricKiwiHOPDataCoordinator
+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 = entry.runtime_data.hop
-        HOP_COORDINATOR
-    ]
 
     _LOGGER.debug("Setting up select entity")
     async_add_entities([ElectricKiwiSelectHOPEntity(hop_coordinator, HOP_SELECT)])

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][
+    account_coordinator = entry.runtime_data.account
-        entry.entry_id
-    ][ACCOUNT_COORDINATOR]
 
     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 = entry.runtime_data.hop
-        HOP_COORDINATOR
-    ]
     entities.extend(
         [
             ElectricKiwiHOPEntity(hop_coordinator, description)