Update sensor entity attributes

Limych · Oct 7, 2024 · 16bb2bb · 16bb2bb
1 parent 8571a64
commit 16bb2bb
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/custom_components/iaquk/sensor.py b/custom_components/iaquk/sensor.py
@@ -6,6 +6,7 @@
 
 from homeassistant.components.sensor import (
     ENTITY_ID_FORMAT,
+    SensorDeviceClass,
     SensorEntity,
     SensorStateClass,
 )
@@ -76,10 +77,7 @@ def __init__(self, controller: IaqukController, sensor_type: str) -> None:
             SensorStateClass.MEASUREMENT if sensor_type == SENSOR_INDEX else None
         )
         self._attr_device_class = (
-            f"{DOMAIN}__level" if sensor_type == SENSOR_LEVEL else None
-        )
-        self._attr_native_unit_of_measurement = (
-            "IAQI" if sensor_type == SENSOR_INDEX else None
+            f"{DOMAIN}__level" if sensor_type == SENSOR_LEVEL else SensorDeviceClass.AQI
         )
         self._attr_icon = ICON_DEFAULT if sensor_type == SENSOR_INDEX else ICON_FAIR
 

diff --git a/tests/test_sensor.py b/tests/test_sensor.py
@@ -2,6 +2,7 @@
 
 from unittest.mock import patch
 
+from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 from pytest_homeassistant_custom_component.common import assert_setup_component
@@ -36,10 +37,11 @@ async def test_entity_initialization(hass: HomeAssistant):
     assert entity.name == "Test Indoor Air Quality Index"
     assert entity.should_poll is True
     assert entity.available is True
+    assert entity.device_class == SensorDeviceClass.AQI
     assert entity.state is None
-    assert entity.state_class == "measurement"
+    assert entity.state_class == SensorStateClass.MEASUREMENT
     assert entity.icon == ICON_DEFAULT
-    assert entity.unit_of_measurement == "IAQI"
+    assert entity.unit_of_measurement is None
     assert entity.extra_state_attributes == expected_attributes
 
     entity = IaqukSensor(controller, SENSOR_LEVEL)
@@ -49,6 +51,7 @@ async def test_entity_initialization(hass: HomeAssistant):
     assert entity.name == "Test Indoor Air Quality Level"
     assert entity.should_poll is True
     assert entity.available is True
+    assert entity.device_class == "iaquk__level"
     assert entity.state is None
     assert entity.state_class is None
     assert entity.icon == ICON_FAIR