Extract Monzo coordinator in separate module (#117034)

This commit is contained in:
Joost Lekkerkerker 2024-05-17 12:40:19 +02:00 committed by GitHub
parent addc4a84ff
commit 098ba125d1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 68 additions and 79 deletions

View file

@ -2,55 +2,35 @@
from __future__ import annotations
from datetime import timedelta
import logging
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.config_entry_oauth2_flow import (
OAuth2Session,
async_get_config_entry_implementation,
)
from .api import AuthenticatedMonzoAPI
from .const import DOMAIN
from .data import MonzoData, MonzoSensorData
from .coordinator import MonzoCoordinator
PLATFORMS: list[Platform] = [Platform.SENSOR]
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up Monzo from a config entry."""
implementation = (
await config_entry_oauth2_flow.async_get_config_entry_implementation(
hass, entry
)
)
implementation = await async_get_config_entry_implementation(hass, entry)
async def async_get_monzo_api_data() -> MonzoSensorData:
monzo_data: MonzoData = hass.data[DOMAIN][entry.entry_id]
accounts = await external_api.user_account.accounts()
pots = await external_api.user_account.pots()
monzo_data.accounts = accounts
monzo_data.pots = pots
return MonzoSensorData(accounts=accounts, pots=pots)
session = OAuth2Session(hass, entry, implementation)
session = config_entry_oauth2_flow.OAuth2Session(hass, entry, implementation)
external_api = AuthenticatedMonzoAPI(async_get_clientsession(hass), session)
external_api = AuthenticatedMonzoAPI(
aiohttp_client.async_get_clientsession(hass), session
)
coordinator = DataUpdateCoordinator(
hass,
logging.getLogger(__name__),
name=DOMAIN,
update_method=async_get_monzo_api_data,
update_interval=timedelta(minutes=1),
)
hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = MonzoData(external_api, coordinator)
coordinator = MonzoCoordinator(hass, external_api)
await coordinator.async_config_entry_first_refresh()
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
return True
@ -58,11 +38,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry."""
data = hass.data[DOMAIN]
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
if unload_ok and entry.entry_id in data:
data.pop(entry.entry_id)
if unload_ok:
hass.data[DOMAIN].pop(entry.entry_id)
return unload_ok

View file

@ -0,0 +1,42 @@
"""The Monzo integration."""
from dataclasses import dataclass
from datetime import timedelta
import logging
from typing import Any
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from .api import AuthenticatedMonzoAPI
from .const import DOMAIN
_LOGGER = logging.getLogger(__name__)
@dataclass
class MonzoData:
"""A dataclass for holding sensor data returned by the DataUpdateCoordinator."""
accounts: list[dict[str, Any]]
pots: list[dict[str, Any]]
class MonzoCoordinator(DataUpdateCoordinator[MonzoData]):
"""Class to manage fetching Monzo data from the API."""
def __init__(self, hass: HomeAssistant, api: AuthenticatedMonzoAPI) -> None:
"""Initialize."""
super().__init__(
hass,
_LOGGER,
name=DOMAIN,
update_interval=timedelta(minutes=1),
)
self.api = api
async def _async_update_data(self) -> MonzoData:
"""Fetch data from Monzo API."""
accounts = await self.api.user_account.accounts()
pots = await self.api.user_account.pots()
return MonzoData(accounts, pots)

View file

@ -1,24 +0,0 @@
"""Dataclass for Monzo data."""
from dataclasses import dataclass, field
from typing import Any
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from .api import AuthenticatedMonzoAPI
@dataclass(kw_only=True)
class MonzoSensorData:
"""A dataclass for holding sensor data returned by the DataUpdateCoordinator."""
accounts: list[dict[str, Any]] = field(default_factory=list)
pots: list[dict[str, Any]] = field(default_factory=list)
@dataclass
class MonzoData(MonzoSensorData):
"""A dataclass for holding data stored in hass.data."""
external_api: AuthenticatedMonzoAPI
coordinator: DataUpdateCoordinator[MonzoSensorData]

View file

@ -6,16 +6,13 @@ from collections.abc import Callable
from typing import Any
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
from homeassistant.helpers.update_coordinator import (
CoordinatorEntity,
DataUpdateCoordinator,
)
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import DOMAIN
from .data import MonzoSensorData
from .coordinator import MonzoCoordinator, MonzoData
class MonzoBaseEntity(CoordinatorEntity[DataUpdateCoordinator[MonzoSensorData]]):
class MonzoBaseEntity(CoordinatorEntity[MonzoCoordinator]):
"""Common base for Monzo entities."""
_attr_attribution = "Data provided by Monzo"
@ -23,10 +20,10 @@ class MonzoBaseEntity(CoordinatorEntity[DataUpdateCoordinator[MonzoSensorData]])
def __init__(
self,
coordinator: DataUpdateCoordinator[MonzoSensorData],
coordinator: MonzoCoordinator,
index: int,
device_model: str,
data_accessor: Callable[[MonzoSensorData], list[dict[str, Any]]],
data_accessor: Callable[[MonzoData], list[dict[str, Any]]],
) -> None:
"""Initialize sensor."""
super().__init__(coordinator)

View file

@ -15,10 +15,10 @@ from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from . import MonzoCoordinator
from .const import DOMAIN
from .data import MonzoSensorData
from .coordinator import MonzoData
from .entity import MonzoBaseEntity
@ -68,7 +68,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Defer sensor setup to the shared sensor module."""
coordinator = hass.data[DOMAIN][config_entry.entry_id].coordinator
coordinator: MonzoCoordinator = hass.data[DOMAIN][config_entry.entry_id]
accounts = [
MonzoSensor(
@ -79,15 +79,13 @@ async def async_setup_entry(
lambda x: x.accounts,
)
for entity_description in ACCOUNT_SENSORS
for index, account in enumerate(
hass.data[DOMAIN][config_entry.entry_id].accounts
)
for index, account in enumerate(coordinator.data.accounts)
]
pots = [
MonzoSensor(coordinator, entity_description, index, MODEL_POT, lambda x: x.pots)
for entity_description in POT_SENSORS
for index, _pot in enumerate(hass.data[DOMAIN][config_entry.entry_id].pots)
for index, _pot in enumerate(coordinator.data.pots)
]
async_add_entities(accounts + pots)
@ -100,11 +98,11 @@ class MonzoSensor(MonzoBaseEntity, SensorEntity):
def __init__(
self,
coordinator: DataUpdateCoordinator[MonzoSensorData],
coordinator: MonzoCoordinator,
entity_description: MonzoSensorEntityDescription,
index: int,
device_model: str,
data_accessor: Callable[[MonzoSensorData], list[dict[str, Any]]],
data_accessor: Callable[[MonzoData], list[dict[str, Any]]],
) -> None:
"""Initialize the sensor."""
super().__init__(coordinator, index, device_model, data_accessor)