Subclass a DataUpdateCoordinator for Ridwell (#85644)

homeassistant/components/ridwell/coordinator.py

@@ -0,0 +1,78 @@
+"""Define a Ridwell coordinator."""
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from typing import cast
+
+from aioridwell.client import async_get_client
+from aioridwell.errors import InvalidCredentialsError, RidwellError
+from aioridwell.model import RidwellAccount, RidwellPickupEvent
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import LOGGER
+
+UPDATE_INTERVAL = timedelta(hours=1)
+
+
+class RidwellDataUpdateCoordinator(
+    DataUpdateCoordinator[dict[str, RidwellPickupEvent]]
+):
+    """Class to manage fetching data from single endpoint."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, *, name: str) -> None:
+        """Initialize global data updater."""
+        # These will be filled in by async_initialize; we give them these defaults to
+        # avoid arduous typing checks down the line:
+        self.accounts: dict[str, RidwellAccount] = {}
+        self.dashboard_url = ""
+        self.user_id = ""
+
+        super().__init__(hass, LOGGER, name=name, update_interval=UPDATE_INTERVAL)
+
+    async def _async_update_data(self) -> dict[str, RidwellPickupEvent]:
+        """Fetch the latest data from the source."""
+        data = {}
+
+        async def async_get_pickups(account: RidwellAccount) -> None:
+            """Get the latest pickups for an account."""
+            data[account.account_id] = await account.async_get_next_pickup_event()
+
+        tasks = [async_get_pickups(account) for account in self.accounts.values()]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        for result in results:
+            if isinstance(result, InvalidCredentialsError):
+                raise ConfigEntryAuthFailed("Invalid username/password") from result
+            if isinstance(result, RidwellError):
+                raise UpdateFailed(result) from result
+
+        return data
+
+    async def async_initialize(self) -> None:
+        """Initialize the coordinator."""
+        session = aiohttp_client.async_get_clientsession(self.hass)
+
+        try:
+            client = await async_get_client(
+                self.config_entry.data[CONF_USERNAME],
+                self.config_entry.data[CONF_PASSWORD],
+                session=session,
+            )
+        except InvalidCredentialsError as err:
+            raise ConfigEntryAuthFailed("Invalid username/password") from err
+        except RidwellError as err:
+            raise ConfigEntryNotReady(err) from err
+
+        self.accounts = await client.async_get_accounts()
+        await self.async_config_entry_first_refresh()
+
+        self.dashboard_url = client.get_dashboard_url()
+        self.user_id = cast(str, client.user_id)