Move WLED sensors to entity descriptions (#58839)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/components/wled/const.py

@@ -17,8 +17,6 @@ ATTR_COLOR_PRIMARY = "color_primary"
 ATTR_DURATION = "duration"
 ATTR_FADE = "fade"
 ATTR_INTENSITY = "intensity"
-ATTR_LED_COUNT = "led_count"
-ATTR_MAX_POWER = "max_power"
 ATTR_ON = "on"
 ATTR_PALETTE = "palette"
 ATTR_PRESET = "preset"

homeassistant/components/wled/sensor.py

@@ -1,10 +1,18 @@
 """Support for WLED sensors."""
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import timedelta
-from typing import Any
+from typing import Callable
 
-from homeassistant.components.sensor import DEVICE_CLASS_CURRENT, SensorEntity
+from wled import Device as WLEDDevice
+
+from homeassistant.components.sensor import (
+    DEVICE_CLASS_CURRENT,
+    STATE_CLASS_MEASUREMENT,
+    SensorEntity,
+    SensorEntityDescription,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     DATA_BYTES,
@@ -17,13 +25,109 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import StateType
 from homeassistant.util.dt import utcnow
 
-from .const import ATTR_LED_COUNT, ATTR_MAX_POWER, DOMAIN
+from .const import DOMAIN
 from .coordinator import WLEDDataUpdateCoordinator
 from .models import WLEDEntity
 
 
+@dataclass
+class WLEDSensorEntityDescriptionMixin:
+    """Mixin for required keys."""
+
+    value_fn: Callable[[WLEDDevice], StateType]
+
+
+@dataclass
+class WLEDSensorEntityDescription(
+    SensorEntityDescription, WLEDSensorEntityDescriptionMixin
+):
+    """Describes WLED sensor entity."""
+
+
+SENSORS: tuple[WLEDSensorEntityDescription, ...] = (
+    WLEDSensorEntityDescription(
+        key="estimated_current",
+        name="Estimated Current",
+        native_unit_of_measurement=ELECTRIC_CURRENT_MILLIAMPERE,
+        device_class=DEVICE_CLASS_CURRENT,
+        state_class=STATE_CLASS_MEASUREMENT,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        value_fn=lambda device: device.info.leds.power,
+    ),
+    WLEDSensorEntityDescription(
+        key="info_leds_count",
+        name="LED Count",
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        value_fn=lambda device: device.info.leds.count,
+    ),
+    WLEDSensorEntityDescription(
+        key="info_leds_max_power",
+        name="Max Current",
+        native_unit_of_measurement=ELECTRIC_CURRENT_MILLIAMPERE,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        device_class=DEVICE_CLASS_CURRENT,
+        value_fn=lambda device: device.info.leds.max_power,
+    ),
+    WLEDSensorEntityDescription(
+        key="uptime",
+        name="Uptime",
+        device_class=DEVICE_CLASS_TIMESTAMP,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+        value_fn=lambda device: (utcnow() - timedelta(seconds=device.info.uptime))
+        .replace(microsecond=0)
+        .isoformat(),
+    ),
+    WLEDSensorEntityDescription(
+        key="free_heap",
+        name="Free Memory",
+        icon="mdi:memory",
+        native_unit_of_measurement=DATA_BYTES,
+        state_class=STATE_CLASS_MEASUREMENT,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+        value_fn=lambda device: device.info.free_heap,
+    ),
+    WLEDSensorEntityDescription(
+        key="wifi_signal",
+        name="Wi-Fi Signal",
+        icon="mdi:wifi",
+        native_unit_of_measurement=PERCENTAGE,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+        value_fn=lambda device: device.info.wifi.signal if device.info.wifi else None,
+    ),
+    WLEDSensorEntityDescription(
+        key="wifi_rssi",
+        name="Wi-Fi RSSI",
+        native_unit_of_measurement=SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
+        device_class=DEVICE_CLASS_SIGNAL_STRENGTH,
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+        value_fn=lambda device: device.info.wifi.rssi if device.info.wifi else None,
+    ),
+    WLEDSensorEntityDescription(
+        key="wifi_channel",
+        name="Wi-Fi Channel",
+        icon="mdi:wifi",
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+        value_fn=lambda device: device.info.wifi.channel if device.info.wifi else None,
+    ),
+    WLEDSensorEntityDescription(
+        key="wifi_bssid",
+        name="Wi-Fi BSSID",
+        icon="mdi:wifi",
+        entity_category=ENTITY_CATEGORY_DIAGNOSTIC,
+        entity_registry_enabled_default=False,
+        value_fn=lambda device: device.info.wifi.bssid if device.info.wifi else None,
+    ),
+)
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: ConfigEntry,
@@ -31,169 +135,28 @@ async def async_setup_entry(
 ) -> None:
     """Set up WLED sensor based on a config entry."""
     coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-
+    async_add_entities(
-    sensors = [
+        WLEDSensorEntity(coordinator, description) for description in SENSORS
-        WLEDEstimatedCurrentSensor(coordinator),
+    )
-        WLEDUptimeSensor(coordinator),
-        WLEDFreeHeapSensor(coordinator),
-        WLEDWifiBSSIDSensor(coordinator),
-        WLEDWifiChannelSensor(coordinator),
-        WLEDWifiRSSISensor(coordinator),
-        WLEDWifiSignalSensor(coordinator),
-    ]
-
-    async_add_entities(sensors)
 
 
-class WLEDEstimatedCurrentSensor(WLEDEntity, SensorEntity):
+class WLEDSensorEntity(WLEDEntity, SensorEntity):
-    """Defines a WLED estimated current sensor."""
+    """Defines a WLED sensor entity."""
 
-    _attr_icon = "mdi:power"
+    entity_description: WLEDSensorEntityDescription
-    _attr_native_unit_of_measurement = ELECTRIC_CURRENT_MILLIAMPERE
-    _attr_device_class = DEVICE_CLASS_CURRENT
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
 
-    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
+    def __init__(
-        """Initialize WLED estimated current sensor."""
+        self,
+        coordinator: WLEDDataUpdateCoordinator,
+        description: WLEDSensorEntityDescription,
+    ) -> None:
+        """Initialize a WLED sensor entity."""
         super().__init__(coordinator=coordinator)
-        self._attr_name = f"{coordinator.data.info.name} Estimated Current"
+        self.entity_description = description
-        self._attr_unique_id = f"{coordinator.data.info.mac_address}_estimated_current"
+        self._attr_name = f"{coordinator.data.info.name} {description.name}"
+        self._attr_unique_id = f"{coordinator.data.info.mac_address}_{description.key}"
 
     @property
-    def extra_state_attributes(self) -> dict[str, Any] | None:
+    def native_value(self) -> StateType:
-        """Return the state attributes of the entity."""
-        return {
-            ATTR_LED_COUNT: self.coordinator.data.info.leds.count,
-            ATTR_MAX_POWER: self.coordinator.data.info.leds.max_power,
-        }
-
-    @property
-    def native_value(self) -> int:
         """Return the state of the sensor."""
-        return self.coordinator.data.info.leds.power
+        return self.entity_description.value_fn(self.coordinator.data)
-
-
-class WLEDUptimeSensor(WLEDEntity, SensorEntity):
-    """Defines a WLED uptime sensor."""
-
-    _attr_device_class = DEVICE_CLASS_TIMESTAMP
-    _attr_entity_registry_enabled_default = False
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
-
-    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
-        """Initialize WLED uptime sensor."""
-        super().__init__(coordinator=coordinator)
-        self._attr_name = f"{coordinator.data.info.name} Uptime"
-        self._attr_unique_id = f"{coordinator.data.info.mac_address}_uptime"
-
-    @property
-    def native_value(self) -> str:
-        """Return the state of the sensor."""
-        uptime = utcnow() - timedelta(seconds=self.coordinator.data.info.uptime)
-        return uptime.replace(microsecond=0).isoformat()
-
-
-class WLEDFreeHeapSensor(WLEDEntity, SensorEntity):
-    """Defines a WLED free heap sensor."""
-
-    _attr_icon = "mdi:memory"
-    _attr_entity_registry_enabled_default = False
-    _attr_native_unit_of_measurement = DATA_BYTES
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
-
-    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
-        """Initialize WLED free heap sensor."""
-        super().__init__(coordinator=coordinator)
-        self._attr_name = f"{coordinator.data.info.name} Free Memory"
-        self._attr_unique_id = f"{coordinator.data.info.mac_address}_free_heap"
-
-    @property
-    def native_value(self) -> int:
-        """Return the state of the sensor."""
-        return self.coordinator.data.info.free_heap
-
-
-class WLEDWifiSignalSensor(WLEDEntity, SensorEntity):
-    """Defines a WLED Wi-Fi signal sensor."""
-
-    _attr_icon = "mdi:wifi"
-    _attr_native_unit_of_measurement = PERCENTAGE
-    _attr_entity_registry_enabled_default = False
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
-
-    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
-        """Initialize WLED Wi-Fi signal sensor."""
-        super().__init__(coordinator=coordinator)
-        self._attr_name = f"{coordinator.data.info.name} Wi-Fi Signal"
-        self._attr_unique_id = f"{coordinator.data.info.mac_address}_wifi_signal"
-
-    @property
-    def native_value(self) -> int | None:
-        """Return the state of the sensor."""
-        if not self.coordinator.data.info.wifi:
-            return None
-        return self.coordinator.data.info.wifi.signal
-
-
-class WLEDWifiRSSISensor(WLEDEntity, SensorEntity):
-    """Defines a WLED Wi-Fi RSSI sensor."""
-
-    _attr_device_class = DEVICE_CLASS_SIGNAL_STRENGTH
-    _attr_native_unit_of_measurement = SIGNAL_STRENGTH_DECIBELS_MILLIWATT
-    _attr_entity_registry_enabled_default = False
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
-
-    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
-        """Initialize WLED Wi-Fi RSSI sensor."""
-        super().__init__(coordinator=coordinator)
-        self._attr_name = f"{coordinator.data.info.name} Wi-Fi RSSI"
-        self._attr_unique_id = f"{coordinator.data.info.mac_address}_wifi_rssi"
-
-    @property
-    def native_value(self) -> int | None:
-        """Return the state of the sensor."""
-        if not self.coordinator.data.info.wifi:
-            return None
-        return self.coordinator.data.info.wifi.rssi
-
-
-class WLEDWifiChannelSensor(WLEDEntity, SensorEntity):
-    """Defines a WLED Wi-Fi Channel sensor."""
-
-    _attr_icon = "mdi:wifi"
-    _attr_entity_registry_enabled_default = False
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
-
-    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
-        """Initialize WLED Wi-Fi Channel sensor."""
-        super().__init__(coordinator=coordinator)
-        self._attr_name = f"{coordinator.data.info.name} Wi-Fi Channel"
-        self._attr_unique_id = f"{coordinator.data.info.mac_address}_wifi_channel"
-
-    @property
-    def native_value(self) -> int | None:
-        """Return the state of the sensor."""
-        if not self.coordinator.data.info.wifi:
-            return None
-        return self.coordinator.data.info.wifi.channel
-
-
-class WLEDWifiBSSIDSensor(WLEDEntity, SensorEntity):
-    """Defines a WLED Wi-Fi BSSID sensor."""
-
-    _attr_icon = "mdi:wifi"
-    _attr_entity_registry_enabled_default = False
-    _attr_entity_category = ENTITY_CATEGORY_DIAGNOSTIC
-
-    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
-        """Initialize WLED Wi-Fi BSSID sensor."""
-        super().__init__(coordinator=coordinator)
-        self._attr_name = f"{coordinator.data.info.name} Wi-Fi BSSID"
-        self._attr_unique_id = f"{coordinator.data.info.mac_address}_wifi_bssid"
-
-    @property
-    def native_value(self) -> str | None:
-        """Return the state of the sensor."""
-        if not self.coordinator.data.info.wifi:
-            return None
-        return self.coordinator.data.info.wifi.bssid

tests/components/wled/test_sensor.py

@@ -10,7 +10,7 @@ from homeassistant.components.sensor import (
     DEVICE_CLASS_TIMESTAMP,
     DOMAIN as SENSOR_DOMAIN,
 )
-from homeassistant.components.wled.const import ATTR_LED_COUNT, ATTR_MAX_POWER, DOMAIN
+from homeassistant.components.wled.const import DOMAIN
 from homeassistant.const import (
     ATTR_DEVICE_CLASS,
     ATTR_ICON,
@@ -95,9 +95,6 @@ async def test_sensors(
 
     state = hass.states.get("sensor.wled_rgb_light_estimated_current")
     assert state
-    assert state.attributes.get(ATTR_ICON) == "mdi:power"
-    assert state.attributes.get(ATTR_LED_COUNT) == 30
-    assert state.attributes.get(ATTR_MAX_POWER) == 850
     assert (
         state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == ELECTRIC_CURRENT_MILLIAMPERE
     )