tries to fix issue where entities have hyphens. also use coordinator to keep track of device availability

Author: zacs

custom_components/wattbox/__init__.py

@@ -24,7 +24,7 @@
     CONF_USERNAME,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import PlatformNotReady
+from homeassistant.exceptions import ConfigEntryNotReady, PlatformNotReady
 from homeassistant.helpers import discovery
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.event import async_track_time_interval
@@ -48,6 +48,7 @@
     STARTUP,
     TOPIC_UPDATE,
 )
+from .coordinator import WattBoxCoordinator
 
 REQUIREMENTS: Final[list[str]] = ["pywattbox>=0.7.2"]
 
@@ -203,18 +204,23 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         wattbox = await _async_create_wattbox(hass, host, port, username, password)
     except Exception as error:
         _LOGGER.error("Error creating WattBox instance: %s", error)
-        raise PlatformNotReady from error
+        raise ConfigEntryNotReady from error
 
     hass.data[DOMAIN_DATA][name] = wattbox
 
+    # Create coordinator for polling and availability tracking
+    coordinator = WattBoxCoordinator(hass, wattbox, name, scan_interval)
+    try:
+        await coordinator.async_config_entry_first_refresh()
+    except ConfigEntryNotReady:
+        raise
+
+    hass.data.setdefault(DOMAIN, {})
+    hass.data[DOMAIN][name] = coordinator
+
     # Forward entry setup to platforms
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
-    # Use the scan interval to trigger updates
-    async_track_time_interval(
-        hass, partial(update_data, hass=hass, name=name), scan_interval
-    )
-
     return True
 
 
@@ -236,5 +242,8 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         # Remove the wattbox from data
         if name in hass.data[DOMAIN_DATA]:
             del hass.data[DOMAIN_DATA][name]
+        # Remove the coordinator
+        if DOMAIN in hass.data and name in hass.data[DOMAIN]:
+            del hass.data[DOMAIN][name]
 
     return unload_ok

custom_components/wattbox/coordinator.py

@@ -0,0 +1,52 @@
+"""DataUpdateCoordinator for WattBox."""
+
+import logging
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from pywattbox.base import BaseWattBox
+
+from datetime import timedelta
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class WattBoxCoordinator(DataUpdateCoordinator[BaseWattBox]):
+    """Coordinator to manage fetching WattBox data from the device.
+
+    When the device is unreachable, UpdateFailed is raised, which causes
+    all entities using this coordinator to automatically become unavailable.
+    When the device comes back, entities automatically become available again.
+    """
+
+    wattbox: BaseWattBox
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        wattbox: BaseWattBox,
+        name: str,
+        scan_interval: timedelta,
+    ) -> None:
+        """Initialize the coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=f"WattBox {name}",
+            update_interval=scan_interval,
+        )
+        self.wattbox = wattbox
+
+    async def _async_update_data(self) -> BaseWattBox:
+        """Fetch data from the WattBox device.
+
+        Raises UpdateFailed on any communication error, which signals
+        to all listening entities that the device is unavailable.
+        """
+        try:
+            await self.wattbox.async_update()
+            return self.wattbox
+        except Exception as err:
+            raise UpdateFailed(
+                f"Error communicating with WattBox: {err}"
+            ) from err

custom_components/wattbox/entity.py

@@ -10,6 +10,7 @@
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import DeviceInfo, Entity
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from pywattbox.base import BaseWattBox
 
 from .const import DOMAIN, DOMAIN_DATA, TOPIC_UPDATE
@@ -43,12 +44,18 @@ class WattBoxEntity(Entity):
     """WattBox Entity class."""
 
     _wattbox: BaseWattBox
+    _coordinator: DataUpdateCoordinator | None
     _async_unsub_dispatcher_connect: Callable
     _attr_should_poll: Literal[False] = False
 
     def __init__(self, hass: HomeAssistant, name: str, *_args: Any) -> None:
         self.hass = hass
         self._wattbox = self.hass.data[DOMAIN_DATA][name]
+
+        # Use coordinator if available (config entry path), otherwise
+        # fall back to dispatcher pattern (legacy YAML path).
+        self._coordinator = self.hass.data.get(DOMAIN, {}).get(name)
+
         self.topic: str = TOPIC_UPDATE.format(DOMAIN, name)
         self._attr_extra_state_attributes: dict[str, Any] = {}
 
@@ -75,6 +82,19 @@ def __init__(self, hass: HomeAssistant, name: str, *_args: Any) -> None:
 
         self._attr_device_info = device_info
 
+    @property
+    def available(self) -> bool:
+        """Return True if entity is available.
+
+        When using a coordinator (config entry path), availability is
+        determined by whether the last data update succeeded. This causes
+        all entities to become unavailable when the WattBox device is
+        unreachable, and available again when it comes back.
+        """
+        if self._coordinator is not None:
+            return self._coordinator.last_update_success
+        return True
+
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""
 
@@ -83,11 +103,22 @@ def update() -> None:
             """Update the state."""
             self.async_schedule_update_ha_state(True)
 
-        self._async_unsub_dispatcher_connect = async_dispatcher_connect(
-            self.hass, self.topic, update
-        )
+        if self._coordinator is not None:
+            # Config entry path: listen to coordinator updates
+            self.async_on_remove(
+                self._coordinator.async_add_listener(update)
+            )
+        else:
+            # Legacy YAML path: listen to dispatcher updates
+            self._async_unsub_dispatcher_connect = async_dispatcher_connect(
+                self.hass, self.topic, update
+            )
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect dispatcher listener when removed."""
-        if hasattr(self, "_async_unsub_dispatcher_connect"):
+        # Coordinator listeners are cleaned up via async_on_remove.
+        # Only manually disconnect the dispatcher listener (YAML path).
+        if self._coordinator is None and hasattr(
+            self, "_async_unsub_dispatcher_connect"
+        ):
             self._async_unsub_dispatcher_connect()

custom_components/wattbox/sensor.py

@@ -11,6 +11,7 @@
 from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.util import slugify
 
 from .const import DOMAIN, DOMAIN_DATA, SENSOR_TYPES
 from .entity import WattBoxEntity
@@ -26,7 +27,7 @@ async def async_setup_entry(
     """Set up WattBox sensors from a config entry."""
     try:
         conf_name: str = entry.data[CONF_NAME]
-        clean_name = conf_name.replace(" ", "_").lower()
+        clean_name = slugify(conf_name)
         entities: list[WattBoxSensor | WattBoxIntegrationSensor] = []
 
         # Get available resources from entry data or use all sensor types
@@ -75,7 +76,7 @@ async def async_setup_platform(
     """Setup sensor platform (legacy YAML support)."""
     try:
         conf_name: str = discovery_info[CONF_NAME]
-        clean_name = conf_name.replace(" ", "_").lower()
+        clean_name = slugify(conf_name)
         entities: list[WattBoxSensor | WattBoxIntegrationSensor] = []
 
         resource: str