Utilize available state for coordinator

Author: brefra

custom_components/plugwise_usb/coordinator.py

@@ -1,19 +1,18 @@
 """DataUpdateCoordinator for Plugwise USB-Stick."""
 
+import asyncio
 from datetime import timedelta
 import logging
 from typing import Any
 
-import async_timeout
-
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.update_coordinator import (
+    ConfigEntryError,
     DataUpdateCoordinator,
     UpdateFailed,
 )
-from plugwise_usb.api import NodeFeature
-from plugwise_usb.exceptions import NodeError, NodeTimeout, StickError
+from plugwise_usb.api import NodeFeature, NodeInfo
+from plugwise_usb.exceptions import NodeError, NodeTimeout, StickError, StickTimeout
 from plugwise_usb.nodes import PlugwiseNode
 
 _LOGGER = logging.getLogger(__name__)
@@ -52,32 +51,35 @@ def __init__(
                 always_update=True,
             )
 
+    async def get_node_info(self) -> NodeInfo:
+        """Return node information."""
+        node_info: NodeInfo | None = None
+        if (node_info := await self.node.node_info_update()) is None:
+            raise UpdateFailed(f"Failed to pull node information from Plugwise node: {self.node.mac}")
+        if node_info.battery_powered:
+            self.always_update = False
+        return node_info
+
     async def async_node_update(self) -> dict[NodeFeature, Any]:
         """Request status update for Plugwise Node."""
+        states: dict[NodeFeature, Any] = {}
+        if not self._initial_update_done:
+            _LOGGER.debug("Initial coordinator update for %s", self.node.mac)
+            states[NodeFeature.INFO] = await self.get_node_info()
+            self._initial_update_done = True
+            return states
+
+        features = tuple(self.async_contexts())
+        _LOGGER.debug("Coordinator update for %s, context=%s", self.node.mac, features)
         try:
-            async with async_timeout.timeout(30):
-                if not self._initial_update_done:
-                    _LOGGER.debug(
-                        "Initial coordinator update for %s", self.node.mac
-                    )
-                    self._initial_update_done = await self.node.node_info_update()
-                    states: dict[NodeFeature, Any] = {}
-                    states[NodeFeature.INFO] = self.node.node_info
-                    if self.node.node_info.battery_powered:
-                        self.always_update = False
-                    return states
-                features = set(self.async_contexts())
-                _LOGGER.debug(
-                    "Coordinator update for %s, context=%s",
-                    self.node.mac,
-                    str(features),
-                )
-                return await self.node.get_state(features)
+            states = await self.node.get_state(features)
         except StickError as err:
-            raise ConfigEntryNotReady from err
-        except NodeTimeout as err:
-            raise TimeoutError from err
+            raise ConfigEntryError from err
+        except (StickTimeout, NodeTimeout) as err:
+            raise asyncio.TimeoutError from err
         except NodeError as err:
-            raise UpdateFailed(
-                f"Failed to pull data from Plugwise node: {err}"
-            ) from err
+            raise UpdateFailed from err
+
+        if not states[NodeFeature.AVAILABLE]:
+            raise UpdateFailed(f"Plugwise node {self.node.mac} is not on-line.")
+        return states