From 098ba125d1e6ab6de7c480b47be343a130e09591 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Fri, 17 May 2024 12:40:19 +0200 Subject: [PATCH] Extract Monzo coordinator in separate module (#117034) --- homeassistant/components/monzo/__init__.py | 52 +++++-------------- homeassistant/components/monzo/coordinator.py | 42 +++++++++++++++ homeassistant/components/monzo/data.py | 24 --------- homeassistant/components/monzo/entity.py | 13 ++--- homeassistant/components/monzo/sensor.py | 16 +++--- 5 files changed, 68 insertions(+), 79 deletions(-) create mode 100644 homeassistant/components/monzo/coordinator.py delete mode 100644 homeassistant/components/monzo/data.py diff --git a/homeassistant/components/monzo/__init__.py b/homeassistant/components/monzo/__init__.py index 93fef56957e..a88082b2ce6 100644 --- a/homeassistant/components/monzo/__init__.py +++ b/homeassistant/components/monzo/__init__.py @@ -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 diff --git a/homeassistant/components/monzo/coordinator.py b/homeassistant/components/monzo/coordinator.py new file mode 100644 index 00000000000..67fff38c4f8 --- /dev/null +++ b/homeassistant/components/monzo/coordinator.py @@ -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) diff --git a/homeassistant/components/monzo/data.py b/homeassistant/components/monzo/data.py deleted file mode 100644 index c4dd2564c21..00000000000 --- a/homeassistant/components/monzo/data.py +++ /dev/null @@ -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] diff --git a/homeassistant/components/monzo/entity.py b/homeassistant/components/monzo/entity.py index 043c06eece0..bf83e3a9bfb 100644 --- a/homeassistant/components/monzo/entity.py +++ b/homeassistant/components/monzo/entity.py @@ -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) diff --git a/homeassistant/components/monzo/sensor.py b/homeassistant/components/monzo/sensor.py index be13608ca3b..41b97d90452 100644 --- a/homeassistant/components/monzo/sensor.py +++ b/homeassistant/components/monzo/sensor.py @@ -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)