"""Define a Notion data coordinator.""" import asyncio from dataclasses import dataclass, field from datetime import timedelta from typing import Any from aionotion.bridge.models import Bridge from aionotion.client import Client from aionotion.errors import InvalidCredentialsError, NotionError from aionotion.listener.models import Listener from aionotion.sensor.models import Sensor from aionotion.user.models import UserPreferences from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_USERNAME from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryAuthFailed from homeassistant.helpers import device_registry as dr from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import DOMAIN, LOGGER DATA_BRIDGES = "bridges" DATA_LISTENERS = "listeners" DATA_SENSORS = "sensors" DATA_USER_PREFERENCES = "user_preferences" DEFAULT_SCAN_INTERVAL = timedelta(minutes=1) @callback def _async_register_new_bridge( hass: HomeAssistant, entry: ConfigEntry, bridge: Bridge ) -> None: """Register a new bridge.""" if name := bridge.name: bridge_name = name.capitalize() else: bridge_name = str(bridge.id) device_registry = dr.async_get(hass) device_registry.async_get_or_create( config_entry_id=entry.entry_id, identifiers={(DOMAIN, bridge.hardware_id)}, manufacturer="Silicon Labs", model=str(bridge.hardware_revision), name=bridge_name, sw_version=bridge.firmware_version.wifi, ) @dataclass class NotionData: """Define a manager class for Notion data.""" hass: HomeAssistant entry: ConfigEntry # Define a dict of bridges, indexed by bridge ID (an integer): bridges: dict[int, Bridge] = field(default_factory=dict) # Define a dict of listeners, indexed by listener UUID (a string): listeners: dict[str, Listener] = field(default_factory=dict) # Define a dict of sensors, indexed by sensor UUID (a string): sensors: dict[str, Sensor] = field(default_factory=dict) # Define a user preferences response object: user_preferences: UserPreferences | None = field(default=None) def update_bridges(self, bridges: list[Bridge]) -> None: """Update the bridges.""" for bridge in bridges: # If a new bridge is discovered, register it: if bridge.id not in self.bridges: _async_register_new_bridge(self.hass, self.entry, bridge) self.bridges[bridge.id] = bridge def update_listeners(self, listeners: list[Listener]) -> None: """Update the listeners.""" self.listeners = {listener.id: listener for listener in listeners} def update_sensors(self, sensors: list[Sensor]) -> None: """Update the sensors.""" self.sensors = {sensor.uuid: sensor for sensor in sensors} def update_user_preferences(self, user_preferences: UserPreferences) -> None: """Update the user preferences.""" self.user_preferences = user_preferences def asdict(self) -> dict[str, Any]: """Represent this dataclass (and its Pydantic contents) as a dict.""" data: dict[str, Any] = { DATA_BRIDGES: [item.to_dict() for item in self.bridges.values()], DATA_LISTENERS: [item.to_dict() for item in self.listeners.values()], DATA_SENSORS: [item.to_dict() for item in self.sensors.values()], } if self.user_preferences: data[DATA_USER_PREFERENCES] = self.user_preferences.to_dict() return data class NotionDataUpdateCoordinator(DataUpdateCoordinator[NotionData]): """Define a Notion data coordinator.""" config_entry: ConfigEntry def __init__( self, hass: HomeAssistant, *, entry: ConfigEntry, client: Client, ) -> None: """Initialize.""" super().__init__( hass, LOGGER, name=entry.data[CONF_USERNAME], update_interval=DEFAULT_SCAN_INTERVAL, ) self._client = client self._entry = entry async def _async_update_data(self) -> NotionData: """Fetch data from Notion.""" data = NotionData(hass=self.hass, entry=self._entry) try: async with asyncio.TaskGroup() as tg: bridges = tg.create_task(self._client.bridge.async_all()) listeners = tg.create_task(self._client.listener.async_all()) sensors = tg.create_task(self._client.sensor.async_all()) user_preferences = tg.create_task(self._client.user.async_preferences()) except BaseExceptionGroup as err: result = err.exceptions[0] if isinstance(result, InvalidCredentialsError): raise ConfigEntryAuthFailed( "Invalid username and/or password" ) from result if isinstance(result, NotionError): raise UpdateFailed( f"There was a Notion error while updating: {result}" ) from result if isinstance(result, Exception): LOGGER.debug( "There was an unknown error while updating: %s", result, exc_info=result, ) raise UpdateFailed( f"There was an unknown error while updating: {result}" ) from result if isinstance(result, BaseException): raise result from None data.update_bridges(bridges.result()) data.update_listeners(listeners.result()) data.update_sensors(sensors.result()) data.update_user_preferences(user_preferences.result()) return data