From 20647af5aeebcb686cd34ef489c65fb0723eb0aa Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Tue, 12 Mar 2024 14:51:13 -0600 Subject: [PATCH] Move Notion coordinator to its own module (#112756) Co-authored-by: Cretezy --- .coveragerc | 1 + homeassistant/components/notion/__init__.py | 140 +-------------- .../components/notion/binary_sensor.py | 4 +- .../components/notion/coordinator.py | 164 ++++++++++++++++++ .../components/notion/diagnostics.py | 5 +- homeassistant/components/notion/sensor.py | 4 +- 6 files changed, 176 insertions(+), 142 deletions(-) create mode 100644 homeassistant/components/notion/coordinator.py diff --git a/.coveragerc b/.coveragerc index c3fd8ba71ff..a60f216af1b 100644 --- a/.coveragerc +++ b/.coveragerc @@ -891,6 +891,7 @@ omit = homeassistant/components/notify_events/notify.py homeassistant/components/notion/__init__.py homeassistant/components/notion/binary_sensor.py + homeassistant/components/notion/coordinator.py homeassistant/components/notion/sensor.py homeassistant/components/notion/util.py homeassistant/components/nsw_fuel_station/sensor.py diff --git a/homeassistant/components/notion/__init__.py b/homeassistant/components/notion/__init__.py index 9b62eb840c0..ca45e3a6d16 100644 --- a/homeassistant/components/notion/__init__.py +++ b/homeassistant/components/notion/__init__.py @@ -2,8 +2,6 @@ from __future__ import annotations -import asyncio -from dataclasses import dataclass, field from datetime import timedelta from typing import Any from uuid import UUID @@ -11,8 +9,6 @@ from uuid import UUID from aionotion.bridge.models import Bridge from aionotion.errors import InvalidCredentialsError, NotionError from aionotion.listener.models import Listener, ListenerKind -from aionotion.sensor.models import Sensor -from aionotion.user.models import UserPreferences from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform @@ -25,11 +21,7 @@ from homeassistant.helpers import ( ) from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity import EntityDescription -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, - UpdateFailed, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import ( CONF_REFRESH_TOKEN, @@ -47,6 +39,7 @@ from .const import ( SENSOR_TEMPERATURE, SENSOR_WINDOW_HINGED, ) +from .coordinator import NotionDataUpdateCoordinator from .util import async_get_client_with_credentials, async_get_client_with_refresh_token PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR] @@ -54,11 +47,6 @@ PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR] ATTR_SYSTEM_MODE = "system_mode" ATTR_SYSTEM_NAME = "system_name" -DATA_BRIDGES = "bridges" -DATA_LISTENERS = "listeners" -DATA_SENSORS = "sensors" -DATA_USER_PREFERENCES = "user_preferences" - DEFAULT_SCAN_INTERVAL = timedelta(minutes=1) CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False) @@ -88,57 +76,6 @@ def is_uuid(value: str) -> bool: return True -@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 - - async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Notion as a config entry.""" entry_updates: dict[str, Any] = {"data": {**entry.data}} @@ -188,51 +125,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: # Create a callback to save the refresh token when it changes: entry.async_on_unload(client.add_refresh_token_callback(async_save_refresh_token)) - async def async_update() -> NotionData: - """Get the latest data from the Notion API.""" - data = NotionData(hass=hass, entry=entry) - - try: - async with asyncio.TaskGroup() as tg: - bridges = tg.create_task(client.bridge.async_all()) - listeners = tg.create_task(client.listener.async_all()) - sensors = tg.create_task(client.sensor.async_all()) - user_preferences = tg.create_task(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 - - coordinator = DataUpdateCoordinator( - hass, - LOGGER, - name=entry.data[CONF_USERNAME], - update_interval=DEFAULT_SCAN_INTERVAL, - update_method=async_update, - ) + coordinator = NotionDataUpdateCoordinator(hass, entry=entry, client=client) await coordinator.async_config_entry_first_refresh() hass.data.setdefault(DOMAIN, {}) @@ -282,39 +175,17 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return unload_ok -@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, - ) - - -class NotionEntity(CoordinatorEntity[DataUpdateCoordinator[NotionData]]): +class NotionEntity(CoordinatorEntity[NotionDataUpdateCoordinator]): """Define a base Notion entity.""" _attr_has_entity_name = True def __init__( self, - coordinator: DataUpdateCoordinator[NotionData], + coordinator: NotionDataUpdateCoordinator, listener_id: str, sensor_id: str, bridge_id: int, - system_id: str, description: EntityDescription, ) -> None: """Initialize the entity.""" @@ -338,7 +209,6 @@ class NotionEntity(CoordinatorEntity[DataUpdateCoordinator[NotionData]]): self._bridge_id = bridge_id self._listener_id = listener_id self._sensor_id = sensor_id - self._system_id = system_id self.entity_description = description @property diff --git a/homeassistant/components/notion/binary_sensor.py b/homeassistant/components/notion/binary_sensor.py index 0c56e4fc33e..da50a809689 100644 --- a/homeassistant/components/notion/binary_sensor.py +++ b/homeassistant/components/notion/binary_sensor.py @@ -31,6 +31,7 @@ from .const import ( SENSOR_SMOKE_CO, SENSOR_WINDOW_HINGED, ) +from .coordinator import NotionDataUpdateCoordinator from .model import NotionEntityDescription @@ -110,7 +111,7 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up Notion sensors based on a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator: NotionDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] async_add_entities( [ @@ -119,7 +120,6 @@ async def async_setup_entry( listener_id, sensor.uuid, sensor.bridge.id, - sensor.system_id, description, ) for listener_id, listener in coordinator.data.listeners.items() diff --git a/homeassistant/components/notion/coordinator.py b/homeassistant/components/notion/coordinator.py new file mode 100644 index 00000000000..c3fd23abc84 --- /dev/null +++ b/homeassistant/components/notion/coordinator.py @@ -0,0 +1,164 @@ +"""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 diff --git a/homeassistant/components/notion/diagnostics.py b/homeassistant/components/notion/diagnostics.py index 77cf9e2a90f..424e5f7d0ac 100644 --- a/homeassistant/components/notion/diagnostics.py +++ b/homeassistant/components/notion/diagnostics.py @@ -8,10 +8,9 @@ from homeassistant.components.diagnostics import async_redact_data from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_EMAIL, CONF_UNIQUE_ID, CONF_USERNAME from homeassistant.core import HomeAssistant -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import NotionData from .const import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN +from .coordinator import NotionDataUpdateCoordinator CONF_DEVICE_KEY = "device_key" CONF_HARDWARE_ID = "hardware_id" @@ -38,7 +37,7 @@ async def async_get_config_entry_diagnostics( hass: HomeAssistant, entry: ConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: DataUpdateCoordinator[NotionData] = hass.data[DOMAIN][entry.entry_id] + coordinator: NotionDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] return async_redact_data( { diff --git a/homeassistant/components/notion/sensor.py b/homeassistant/components/notion/sensor.py index 9bb8e3b4bbf..d12dabbbc33 100644 --- a/homeassistant/components/notion/sensor.py +++ b/homeassistant/components/notion/sensor.py @@ -17,6 +17,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from . import NotionEntity from .const import DOMAIN, SENSOR_MOLD, SENSOR_TEMPERATURE +from .coordinator import NotionDataUpdateCoordinator from .model import NotionEntityDescription @@ -45,7 +46,7 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up Notion sensors based on a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator: NotionDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] async_add_entities( [ @@ -54,7 +55,6 @@ async def async_setup_entry( listener_id, sensor.uuid, sensor.bridge.id, - sensor.system_id, description, ) for listener_id, listener in coordinator.data.listeners.items()