Move Notion coordinator to its own module (#112756)

Co-authored-by: Cretezy <Cretezy@users.noreply.github.com>

.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

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 (
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-    UpdateFailed,
-)
 
 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:
+    coordinator = NotionDataUpdateCoordinator(hass, entry=entry, client=client)
-        """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,
-    )
 
     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
+class NotionEntity(CoordinatorEntity[NotionDataUpdateCoordinator]):
-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]]):
     """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

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()

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

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(
         {

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()