Add sensor platform to WLED integration (#28632)

* Add sensor platform to WLED integration

* Process review comments

homeassistant/components/wled/__init__.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Union
 from wled import WLED, 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
@@ -34,6 +35,7 @@ from .const import (
 )
 
 SCAN_INTERVAL = timedelta(seconds=5)
+WLED_COMPONENTS = (LIGHT_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -60,7 +62,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     hass.data[DOMAIN][entry.entry_id] = {DATA_WLED_CLIENT: wled}
 
     # Set up all platforms for this device/entry.
-    for component in LIGHT_DOMAIN, SWITCH_DOMAIN:
+    for component in WLED_COMPONENTS:
         hass.async_create_task(
             hass.config_entries.async_forward_entry_setup(entry, component)
         )
@@ -93,8 +95,10 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # Unload entities for this entry/device.
     await asyncio.gather(
-        hass.config_entries.async_forward_entry_unload(entry, LIGHT_DOMAIN),
+        *(
-        hass.config_entries.async_forward_entry_unload(entry, SWITCH_DOMAIN),
+            hass.config_entries.async_forward_entry_unload(entry, component)
+            for component in WLED_COMPONENTS
+        )
     )
 
     # Cleanup

homeassistant/components/wled/const.py

@@ -14,7 +14,9 @@ ATTR_DURATION = "duration"
 ATTR_FADE = "fade"
 ATTR_IDENTIFIERS = "identifiers"
 ATTR_INTENSITY = "intensity"
+ATTR_LED_COUNT = "led_count"
 ATTR_MANUFACTURER = "manufacturer"
+ATTR_MAX_POWER = "max_power"
 ATTR_MODEL = "model"
 ATTR_ON = "on"
 ATTR_PALETTE = "palette"
@@ -25,3 +27,7 @@ ATTR_SOFTWARE_VERSION = "sw_version"
 ATTR_SPEED = "speed"
 ATTR_TARGET_BRIGHTNESS = "target_brightness"
 ATTR_UDP_PORT = "udp_port"
+
+# Units of measurement
+CURRENT_MA = "mA"
+DATA_BYTES = "bytes"

homeassistant/components/wled/sensor.py

@@ -0,0 +1,141 @@
+"""Support for WLED sensors."""
+from datetime import timedelta
+import logging
+from typing import Callable, List, Optional, Union
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import DEVICE_CLASS_TIMESTAMP
+from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.typing import HomeAssistantType
+from homeassistant.util.dt import utcnow
+
+from . import WLED, WLEDDeviceEntity
+from .const import (
+    ATTR_LED_COUNT,
+    ATTR_MAX_POWER,
+    CURRENT_MA,
+    DATA_BYTES,
+    DATA_WLED_CLIENT,
+    DOMAIN,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(
+    hass: HomeAssistantType,
+    entry: ConfigEntry,
+    async_add_entities: Callable[[List[Entity], bool], None],
+) -> None:
+    """Set up WLED sensor based on a config entry."""
+    wled: WLED = hass.data[DOMAIN][entry.entry_id][DATA_WLED_CLIENT]
+
+    sensors = [
+        WLEDEstimatedCurrentSensor(entry.entry_id, wled),
+        WLEDUptimeSensor(entry.entry_id, wled),
+        WLEDFreeHeapSensor(entry.entry_id, wled),
+    ]
+
+    async_add_entities(sensors, True)
+
+
+class WLEDSensor(WLEDDeviceEntity):
+    """Defines a WLED sensor."""
+
+    def __init__(
+        self,
+        entry_id: str,
+        wled: WLED,
+        name: str,
+        icon: str,
+        unit_of_measurement: str,
+        key: str,
+    ) -> None:
+        """Initialize WLED sensor."""
+        self._state = None
+        self._unit_of_measurement = unit_of_measurement
+        self._key = key
+
+        super().__init__(entry_id, wled, name, icon)
+
+    @property
+    def unique_id(self) -> str:
+        """Return the unique ID for this sensor."""
+        return f"{self.wled.device.info.mac_address}_{self._key}"
+
+    @property
+    def state(self) -> Union[None, str, int, float]:
+        """Return the state of the sensor."""
+        return self._state
+
+    @property
+    def unit_of_measurement(self) -> str:
+        """Return the unit this state is expressed in."""
+        return self._unit_of_measurement
+
+
+class WLEDEstimatedCurrentSensor(WLEDSensor):
+    """Defines a WLED estimated current sensor."""
+
+    def __init__(self, entry_id: str, wled: WLED) -> None:
+        """Initialize WLED estimated current sensor."""
+        super().__init__(
+            entry_id,
+            wled,
+            f"{wled.device.info.name} Estimated Current",
+            "mdi:power",
+            CURRENT_MA,
+            "estimated_current",
+        )
+
+    async def _wled_update(self) -> None:
+        """Update WLED entity."""
+        self._state = self.wled.device.info.leds.power
+        self._attributes = {
+            ATTR_LED_COUNT: self.wled.device.info.leds.count,
+            ATTR_MAX_POWER: self.wled.device.info.leds.max_power,
+        }
+
+
+class WLEDUptimeSensor(WLEDSensor):
+    """Defines a WLED uptime sensor."""
+
+    def __init__(self, entry_id: str, wled: WLED) -> None:
+        """Initialize WLED uptime sensor."""
+        super().__init__(
+            entry_id,
+            wled,
+            f"{wled.device.info.name} Uptime",
+            "mdi:clock-outline",
+            None,
+            "uptime",
+        )
+
+    @property
+    def device_class(self) -> Optional[str]:
+        """Return the class of this sensor."""
+        return DEVICE_CLASS_TIMESTAMP
+
+    async def _wled_update(self) -> None:
+        """Update WLED uptime sensor."""
+        uptime = utcnow() - timedelta(seconds=self.wled.device.info.uptime)
+        self._state = uptime.replace(microsecond=0).isoformat()
+
+
+class WLEDFreeHeapSensor(WLEDSensor):
+    """Defines a WLED free heap sensor."""
+
+    def __init__(self, entry_id: str, wled: WLED) -> None:
+        """Initialize WLED free heap sensor."""
+        super().__init__(
+            entry_id,
+            wled,
+            f"{wled.device.info.name} Free Memory",
+            "mdi:memory",
+            DATA_BYTES,
+            "free_heap",
+        )
+
+    async def _wled_update(self) -> None:
+        """Update WLED uptime sensor."""
+        self._state = self.wled.device.info.free_heap

tests/components/wled/test_sensor.py

@@ -0,0 +1,61 @@
+"""Tests for the WLED sensor platform."""
+from datetime import datetime
+
+from asynctest import patch
+
+from homeassistant.components.wled.const import (
+    ATTR_LED_COUNT,
+    ATTR_MAX_POWER,
+    CURRENT_MA,
+    DATA_BYTES,
+)
+from homeassistant.const import ATTR_ICON, ATTR_UNIT_OF_MEASUREMENT
+from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
+
+from tests.components.wled import init_integration
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+
+async def test_sensors(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the creation and values of the WLED sensors."""
+
+    test_time = datetime(2019, 11, 11, 9, 10, 32, tzinfo=dt_util.UTC)
+    with patch("homeassistant.components.wled.sensor.utcnow", return_value=test_time):
+        await init_integration(hass, aioclient_mock)
+
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
+
+    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) == CURRENT_MA
+    assert state.state == "470"
+
+    entry = entity_registry.async_get("sensor.wled_rgb_light_estimated_current")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_estimated_current"
+
+    state = hass.states.get("sensor.wled_rgb_light_uptime")
+    assert state
+    assert state.attributes.get(ATTR_ICON) == "mdi:clock-outline"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) is None
+    assert state.state == "2019-11-11T09:10:00+00:00"
+
+    entry = entity_registry.async_get("sensor.wled_rgb_light_uptime")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_uptime"
+
+    state = hass.states.get("sensor.wled_rgb_light_free_memory")
+    assert state
+    assert state.attributes.get(ATTR_ICON) == "mdi:memory"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == DATA_BYTES
+    assert state.state == "14600"
+
+    entry = entity_registry.async_get("sensor.wled_rgb_light_free_memory")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_free_heap"