diff --git a/homeassistant/components/wled/sensor.py b/homeassistant/components/wled/sensor.py index b684bdc0977..8cfd9f98f02 100644 --- a/homeassistant/components/wled/sensor.py +++ b/homeassistant/components/wled/sensor.py @@ -3,6 +3,7 @@ from datetime import timedelta import logging from typing import Any, Callable, Dict, List, Optional, Union +from homeassistant.components.sensor import DEVICE_CLASS_CURRENT from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( DATA_BYTES, @@ -105,6 +106,11 @@ class WLEDEstimatedCurrentSensor(WLEDSensor): """Return the state of the sensor.""" return self.coordinator.data.info.leds.power + @property + def device_class(self) -> Optional[str]: + """Return the class of this sensor.""" + return DEVICE_CLASS_CURRENT + class WLEDUptimeSensor(WLEDSensor): """Defines a WLED uptime sensor.""" diff --git a/tests/components/wled/test_sensor.py b/tests/components/wled/test_sensor.py index 66fedfdd274..adbd2e6e1fc 100644 --- a/tests/components/wled/test_sensor.py +++ b/tests/components/wled/test_sensor.py @@ -3,7 +3,10 @@ from datetime import datetime import pytest -from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN +from homeassistant.components.sensor import ( + DEVICE_CLASS_CURRENT, + DOMAIN as SENSOR_DOMAIN, +) from homeassistant.components.wled.const import ( ATTR_LED_COUNT, ATTR_MAX_POWER, @@ -12,6 +15,7 @@ from homeassistant.components.wled.const import ( SIGNAL_DBM, ) from homeassistant.const import ( + ATTR_DEVICE_CLASS, ATTR_ICON, ATTR_UNIT_OF_MEASUREMENT, DATA_BYTES, @@ -94,6 +98,7 @@ async def test_sensors( 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.attributes.get(ATTR_DEVICE_CLASS) == DEVICE_CLASS_CURRENT assert state.state == "470" entry = registry.async_get("sensor.wled_rgb_light_estimated_current")