Extract Monzo coordinator in separate module (#117034)

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

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)

homeassistant/components/monzo/data.py

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

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)

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)