Migrate WLED to use DataUpdateCoordinator

homeassistant/components/wled/__init__.py

@@ -2,35 +2,27 @@
 import asyncio
 from datetime import timedelta
 import logging
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict
 
-from wled import WLED, WLEDConnectionError, WLEDError
+from wled import WLED, Device as WLEDDevice, WLEDConnectionError, WLEDError
 
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_NAME, CONF_HOST
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.dispatcher import (
-    async_dispatcher_connect,
-    async_dispatcher_send,
-)
 from homeassistant.helpers.entity import Entity
-from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.helpers.typing import ConfigType
-from homeassistant.util import dt as dt_util
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import (
     ATTR_IDENTIFIERS,
     ATTR_MANUFACTURER,
     ATTR_MODEL,
     ATTR_SOFTWARE_VERSION,
-    DATA_WLED_CLIENT,
-    DATA_WLED_TIMER,
-    DATA_WLED_UPDATED,
     DOMAIN,
 )
 
@@ -49,22 +41,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up WLED from a config entry."""
 
     # Create WLED instance for this entry
-    session = async_get_clientsession(hass)
-    wled = WLED(entry.data[CONF_HOST], session=session)
+    coordinator = WLEDDataUpdateCoordinator(hass, host=entry.data[CONF_HOST])
+    await coordinator.async_refresh()
 
-    # Ensure we can connect and talk to it
-    try:
-        await wled.update()
-    except WLEDConnectionError as exception:
-        raise ConfigEntryNotReady from exception
+    if not coordinator.last_update_success:
+        raise ConfigEntryNotReady
 
     hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = {DATA_WLED_CLIENT: wled}
+    hass.data[DOMAIN][entry.entry_id] = coordinator
 
     # For backwards compat, set unique ID
     if entry.unique_id is None:
         hass.config_entries.async_update_entry(
-            entry, unique_id=wled.device.info.mac_address
+            entry, unique_id=coordinator.data.info.mac_address
         )
 
     # Set up all platforms for this device/entry.
@@ -73,32 +62,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             hass.config_entries.async_forward_entry_setup(entry, component)
         )
 
-    async def interval_update(now: dt_util.dt.datetime = None) -> None:
-        """Poll WLED device function, dispatches event after update."""
-        try:
-            await wled.update()
-        except WLEDError:
-            _LOGGER.debug("An error occurred while updating WLED", exc_info=True)
-
-        # Even if the update failed, we still send out the event.
-        # To allow entities to make themselves unavailable.
-        async_dispatcher_send(hass, DATA_WLED_UPDATED, entry.entry_id)
-
-    # Schedule update interval
-    hass.data[DOMAIN][entry.entry_id][DATA_WLED_TIMER] = async_track_time_interval(
-        hass, interval_update, SCAN_INTERVAL
-    )
-
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload WLED config entry."""
 
-    # Cancel update timer for this entry/device.
-    cancel_timer = hass.data[DOMAIN][entry.entry_id][DATA_WLED_TIMER]
-    cancel_timer()
-
     # Unload entities for this entry/device.
     await asyncio.gather(
         *(
@@ -115,26 +84,74 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
+def wled_exception_handler(func):
+    """Decorate WLED calls to handle WLED exceptions.
+
+    A decorator that wraps the passed in function, catches WLED errors,
+    and handles the availability of the device in the data coordinator.
+    """
+
+    async def handler(self, *args, **kwargs):
+        try:
+            await func(self, *args, **kwargs)
+            await self.coordinator.async_refresh()
+
+        except WLEDConnectionError as error:
+            _LOGGER.error("Error communicating with API: %s", error)
+            self.coordinator.last_update_success = False
+            self.coordinator.update_listeners()
+
+        except WLEDError as error:
+            _LOGGER.error("Invalid response from API: %s", error)
+
+    return handler
+
+
+class WLEDDataUpdateCoordinator(DataUpdateCoordinator):
+    """Class to manage fetching WLED data from single endpoint."""
+
+    def __init__(
+        self, hass: HomeAssistant, *, host: str,
+    ):
+        """Initialize global WLED data updater."""
+        self.wled = WLED(host, session=async_get_clientsession(hass))
+
+        super().__init__(
+            hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL,
+        )
+
+    def update_listeners(self) -> None:
+        """Call update on all listeners."""
+        for update_callback in self._listeners:
+            update_callback()
+
+    async def _async_update_data(self) -> WLEDDevice:
+        """Fetch data from WLED."""
+        try:
+            return await self.wled.update()
+        except WLEDError as error:
+            raise UpdateFailed(f"Invalid response from API: {error}")
+
+
 class WLEDEntity(Entity):
     """Defines a base WLED entity."""
 
     def __init__(
         self,
+        *,
         entry_id: str,
-        wled: WLED,
+        coordinator: WLEDDataUpdateCoordinator,
         name: str,
         icon: str,
         enabled_default: bool = True,
     ) -> None:
         """Initialize the WLED entity."""
-        self._attributes: Dict[str, Union[str, int, float]] = {}
-        self._available = True
         self._enabled_default = enabled_default
         self._entry_id = entry_id
         self._icon = icon
         self._name = name
         self._unsub_dispatcher = None
-        self.wled = wled
+        self.coordinator = coordinator
 
     @property
     def name(self) -> str:
@@ -149,7 +166,7 @@ def icon(self) -> str:
     @property
     def available(self) -> bool:
         """Return True if entity is available."""
-        return self._available
+        return self.coordinator.last_update_success
 
     @property
     def entity_registry_enabled_default(self) -> bool:
@@ -161,42 +178,17 @@ def should_poll(self) -> bool:
         """Return the polling requirement of the entity."""
         return False
 
-    @property
-    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
-        """Return the state attributes of the entity."""
-        return self._attributes
-
     async def async_added_to_hass(self) -> None:
         """Connect to dispatcher listening for entity data notifications."""
-        self._unsub_dispatcher = async_dispatcher_connect(
-            self.hass, DATA_WLED_UPDATED, self._schedule_immediate_update
-        )
+        self.coordinator.async_add_listener(self.async_write_ha_state)
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect from update signal."""
-        self._unsub_dispatcher()
-
-    @callback
-    def _schedule_immediate_update(self, entry_id: str) -> None:
-        """Schedule an immediate update of the entity."""
-        if entry_id == self._entry_id:
-            self.async_schedule_update_ha_state(True)
+        self.coordinator.async_remove_listener(self.async_write_ha_state)
 
     async def async_update(self) -> None:
         """Update WLED entity."""
-        if not self.enabled:
-            return
-
-        if self.wled.device is None:
-            self._available = False
-            return
-
-        self._available = True
-        await self._wled_update()
-
-    async def _wled_update(self) -> None:
-        """Update WLED entity."""
-        raise NotImplementedError()
+        await self.coordinator.async_request_refresh()
 
 
 class WLEDDeviceEntity(WLEDEntity):
@@ -206,9 +198,9 @@ class WLEDDeviceEntity(WLEDEntity):
     def device_info(self) -> Dict[str, Any]:
         """Return device information about this WLED device."""
         return {
-            ATTR_IDENTIFIERS: {(DOMAIN, self.wled.device.info.mac_address)},
-            ATTR_NAME: self.wled.device.info.name,
-            ATTR_MANUFACTURER: self.wled.device.info.brand,
-            ATTR_MODEL: self.wled.device.info.product,
-            ATTR_SOFTWARE_VERSION: self.wled.device.info.version,
+            ATTR_IDENTIFIERS: {(DOMAIN, self.coordinator.data.info.mac_address)},
+            ATTR_NAME: self.coordinator.data.info.name,
+            ATTR_MANUFACTURER: self.coordinator.data.info.brand,
+            ATTR_MODEL: self.coordinator.data.info.product,
+            ATTR_SOFTWARE_VERSION: self.coordinator.data.info.version,
         }

homeassistant/components/wled/const.py

@@ -3,11 +3,6 @@
 # Integration domain
 DOMAIN = "wled"
 
-# Home Assistant data keys
-DATA_WLED_CLIENT = "wled_client"
-DATA_WLED_TIMER = "wled_timer"
-DATA_WLED_UPDATED = "wled_updated"
-
 # Attributes
 ATTR_COLOR_PRIMARY = "color_primary"
 ATTR_DURATION = "duration"