Add device registry support to group entities under their stations (#7)

feat: add last update sensor to track data freshness closes #1
elboletaire · Jan 2, 2025 · 2d0a6d5 · 2d0a6d5
1 parent c83704b
commit 2d0a6d5
Show file tree

Hide file tree

Showing 7 changed files with 129 additions and 2 deletions.
diff --git a/custom_components/weatherxm/battery.py b/custom_components/weatherxm/battery.py
@@ -1,4 +1,5 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.entity import DeviceInfo
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -7,6 +8,8 @@
     SensorStateClass,
 )
 
+from .const import DOMAIN
+
 BATTERY_LEVEL_MAP = {
     "ok": 100.0,
     "low": 3.0,
@@ -70,3 +73,13 @@ def icon(self):
         if not self._is_active:
             return "mdi:battery-alert"
         return "mdi:battery" if self._bat_state == "ok" else "mdi:battery-alert"
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._device_id)},
+            name=self._alias,
+            manufacturer="WeatherXM",
+            model="Weather Station",
+        )
diff --git a/custom_components/weatherxm/firmware.py b/custom_components/weatherxm/firmware.py
@@ -1,6 +1,9 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.components.sensor import SensorEntity
 
+from .const import DOMAIN
+
 class WeatherXMFirmwareSensor(CoordinatorEntity, SensorEntity):
     def __init__(self, coordinator, device_id, alias, firmware):
         super().__init__(coordinator)
@@ -39,3 +42,14 @@ def extra_state_attributes(self):
     @property
     def icon(self):
         return "mdi:chip"
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._device_id)},
+            name=self._alias,
+            manufacturer="WeatherXM",
+            model="Weather Station",
+            sw_version=self.state
+        )
diff --git a/custom_components/weatherxm/geo_location.py b/custom_components/weatherxm/geo_location.py
@@ -3,7 +3,7 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
-from homeassistant.helpers.entity import generate_entity_id
+from homeassistant.helpers.entity import generate_entity_id, DeviceInfo
 
 from .const import DOMAIN
 from .utils import async_setup_entities_list
@@ -89,3 +89,13 @@ def extra_state_attributes(self):
     @property
     def icon(self):
         return "mdi:map-marker"
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._device_id)},
+            name=self._alias,
+            manufacturer="WeatherXM",
+            model="Weather Station",
+        )
diff --git a/custom_components/weatherxm/last_update.py b/custom_components/weatherxm/last_update.py
@@ -0,0 +1,48 @@
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.components.sensor import (
+    SensorEntity,
+    SensorDeviceClass,
+)
+
+from .const import DOMAIN
+
+class WeatherXMLastUpdateSensor(CoordinatorEntity, SensorEntity):
+    """WeatherXM Last Update Sensor."""
+
+    _attr_device_class = SensorDeviceClass.TIMESTAMP
+
+    def __init__(self, coordinator, device_id, alias):
+        """Initialize."""
+        super().__init__(coordinator)
+        self._device_id = device_id
+        self._alias = alias
+        self._attr_name = f"{alias} Last Update"
+        self._attr_unique_id = f"{alias}_last_update"
+
+    def _get_device_data(self):
+        """Get device data from coordinator."""
+        if not self.coordinator.data:
+            return None
+        for device in self.coordinator.data:
+            if device['id'] == self._device_id:
+                return device
+        return None
+
+    @property
+    def state(self):
+        """Return the state of the sensor."""
+        device = self._get_device_data()
+        if device and device.get('current_weather', {}).get('timestamp'):
+            return device['current_weather']['timestamp']
+        return None
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._device_id)},
+            name=self._alias,
+            manufacturer="WeatherXM",
+            model="Weather Station",
+        )
diff --git a/custom_components/weatherxm/rewards.py b/custom_components/weatherxm/rewards.py
@@ -1,4 +1,5 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.entity import DeviceInfo
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -7,6 +8,8 @@
     SensorStateClass,
 )
 
+from .const import DOMAIN
+
 class WeatherXMRewardsSensor(CoordinatorEntity, SensorEntity):
     def __init__(self, coordinator, device_id, alias, actual_reward, total_rewards):
         super().__init__(coordinator)
@@ -51,6 +54,16 @@ def unit_of_measurement(self):
     def icon(self):
         return "mdi:currency-usd"
 
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._device_id)},
+            name=self._alias,
+            manufacturer="WeatherXM",
+            model="Weather Station",
+        )
+
 class WeatherXMTotalRewardsSensor(CoordinatorEntity, SensorEntity):
     def __init__(self, coordinator, device_id, alias, total_rewards):
         super().__init__(coordinator)
@@ -87,3 +100,13 @@ def unit_of_measurement(self):
     @property
     def icon(self):
         return "mdi:currency-usd"
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._device_id)},
+            name=self._alias,
+            manufacturer="WeatherXM",
+            model="Weather Station",
+        )
diff --git a/custom_components/weatherxm/sensor.py b/custom_components/weatherxm/sensor.py
@@ -6,6 +6,7 @@
 from .battery import WeatherXMBatteryLevelSensor
 from .rewards import WeatherXMRewardsSensor, WeatherXMTotalRewardsSensor
 from .firmware import WeatherXMFirmwareSensor
+from .last_update import WeatherXMLastUpdateSensor
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_entities):
     coordinator = hass.data[DOMAIN][entry.entry_id]['coordinator']
@@ -47,3 +48,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
         firmware=device['attributes']['firmware']
     ))
     async_add_entities(firmware, True)
+
+    # Last update sensors
+    last_update = await async_setup_entities_list(hass, entry, lambda alias, device: WeatherXMLastUpdateSensor(
+        coordinator=coordinator,
+        device_id=device['id'],
+        alias=alias
+    ))
+    async_add_entities(last_update, True)
diff --git a/custom_components/weatherxm/weather.py b/custom_components/weatherxm/weather.py
@@ -6,7 +6,7 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
-from homeassistant.helpers.entity import generate_entity_id
+from homeassistant.helpers.entity import generate_entity_id, DeviceInfo
 from homeassistant.const import (
     UnitOfPrecipitationDepth,
     UnitOfPressure,
@@ -269,3 +269,13 @@ async def async_forecast_daily(self):
 
     async def async_forecast_hourly(self):
         return self.forecast_hourly[:24]
+
+    @property
+    def device_info(self) -> DeviceInfo:
+        """Return device information."""
+        return DeviceInfo(
+            identifiers={(DOMAIN, self._device_id)},
+            name=self._alias,
+            manufacturer="WeatherXM",
+            model="Weather Station",
+        )