Use runtime_data in electric_kiwi (#136699)

This commit is contained in:
epenet 2025-01-28 11:34:57 +01:00 committed by GitHub
parent 164078ac69
commit c7c234c5dd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 56 additions and 35 deletions

View File

@ -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)

View File

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

View File

@ -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.

View File

@ -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)])

View File

@ -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)