Update component to modern HA version

custom_components/iaquk/sensor.py

@@ -1,12 +1,16 @@
 """Sensor platform to calculate IAQ UK index."""
 
 import logging
-from typing import Any, Dict, Optional, Union
+from typing import Any, Mapping, Optional
 
-from homeassistant.components.sensor import ENTITY_ID_FORMAT
+from homeassistant.components.sensor import (
+    ENTITY_ID_FORMAT,
+    STATE_CLASS_MEASUREMENT,
+    SensorEntity,
+)
 from homeassistant.const import CONF_NAME, CONF_SENSORS
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import Entity, async_generate_entity_id
+from homeassistant.helpers.entity import async_generate_entity_id
 from homeassistant.helpers.typing import ConfigType
 
 from .const import (
@@ -53,72 +57,53 @@ async def async_setup_platform(
     async_add_entities(sensors, True)
 
 
-class IaqukSensor(Entity):
+class IaqukSensor(SensorEntity):
     """IAQ UK sensor."""
 
     def __init__(self, controller, sensor_type: str):
         """Initialize sensor."""
         self._controller = controller
         self._sensor_type = sensor_type
-        self._unique_id = f"{self._controller.unique_id}_{self._sensor_type}"
-        self._name = "{} {}".format(self._controller.name, SENSORS[self._sensor_type])
 
         self.entity_id = async_generate_entity_id(
-            ENTITY_ID_FORMAT, self._unique_id, hass=controller.hass
+            ENTITY_ID_FORMAT, self._attr_unique_id, hass=controller.hass
         )
 
+        self._attr_unique_id = f"{controller.unique_id}_{sensor_type}"
+        self._attr_name = "{} {}".format(controller.name, SENSORS[sensor_type])
+        self._attr_state_class = STATE_CLASS_MEASUREMENT
+        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
+        )
+        self._attr_icon = ICON_DEFAULT if sensor_type == SENSOR_INDEX else ICON_FAIR
+
     async def async_added_to_hass(self):
         """Register callbacks."""
         self._controller.async_added_to_hass()
 
     @property
-    def unique_id(self) -> Optional[str]:
-        """Return a unique ID."""
-        return self._unique_id
-
-    @property
-    def name(self) -> Optional[str]:
-        """Return the name of the sensor."""
-        return self._name
-
-    @property
-    def icon(self) -> Optional[str]:
-        """Icon to use in the frontend, if any."""
-        icon = ICON_DEFAULT
-        if self._sensor_type == SENSOR_LEVEL:
-            icon = ICON_FAIR
-            if self.state == LEVEL_EXCELLENT:
-                icon = ICON_EXCELLENT
-            elif self.state == LEVEL_GOOD:
-                icon = ICON_GOOD
-            # Skip for LEVEL_FAIR -- default state
-            elif self.state == LEVEL_POOR:
-                icon = ICON_POOR
-            elif self.state == LEVEL_INADEQUATE:
-                icon = ICON_INADEQUATE
-
-        return icon
-
-    @property
-    def state(self) -> Union[None, str, int, float]:
-        """Return the state of the sensor."""
-        return (
-            self._controller.iaq_index
-            if self._sensor_type == SENSOR_INDEX
-            else self._controller.iaq_level
-        )
-
-    @property
-    def device_class(self) -> Optional[str]:
-        """Return the class of this device, from component DEVICE_CLASSES."""
-        return f"{DOMAIN}__level" if self._sensor_type == SENSOR_LEVEL else None
-
-    @property
-    def unit_of_measurement(self) -> Optional[str]:
-        """Return the unit of measurement of this entity, if any."""
-        return "IAQI" if self._sensor_type == SENSOR_INDEX else None
-
-    @property
-    def state_attributes(self) -> Optional[Dict[str, Any]]:
+    def extra_state_attributes(self) -> Optional[Mapping[str, Any]]:
         """Return the state attributes."""
         return self._controller.state_attributes
+
+    async def async_update(self):
+        """Update sensor state."""
+        if self._sensor_type == SENSOR_INDEX:
+            self._attr_native_value = self._controller.iaq_index
+
+        else:
+            self._attr_native_value = self._controller.iaq_level
+            self._attr_icon = (
+                ICON_EXCELLENT
+                if self.state == LEVEL_EXCELLENT
+                else ICON_GOOD
+                if self.state == LEVEL_GOOD
+                else ICON_POOR
+                if self.state == LEVEL_POOR
+                else ICON_INADEQUATE
+                if self.state == LEVEL_INADEQUATE
+                else ICON_FAIR
+            )

tests/test_sensor.py

@@ -29,6 +29,7 @@ async def test_entity_initialization(hass: HomeAssistant):
     expected_attributes = {"sources_set": 1, "sources_used": 0}
 
     entity = IaqukSensor(controller, SENSOR_INDEX)
+    entity.hass = hass
 
     assert entity.unique_id == "test_iaq_index"
     assert entity.name == "Test Indoor Air Quality Index"
@@ -37,9 +38,10 @@ async def test_entity_initialization(hass: HomeAssistant):
     assert entity.state is None
     assert entity.icon == ICON_DEFAULT
     assert entity.unit_of_measurement == "IAQI"
-    assert entity.state_attributes == expected_attributes
+    assert entity.extra_state_attributes == expected_attributes
 
     entity = IaqukSensor(controller, SENSOR_LEVEL)
+    entity.hass = hass
 
     assert entity.unique_id == "test_iaq_level"
     assert entity.name == "Test Indoor Air Quality Level"
@@ -48,7 +50,7 @@ async def test_entity_initialization(hass: HomeAssistant):
     assert entity.state is None
     assert entity.icon == ICON_FAIR
     assert entity.unit_of_measurement is None
-    assert entity.state_attributes == expected_attributes
+    assert entity.extra_state_attributes == expected_attributes
 
     levels = {
         LEVEL_EXCELLENT: ICON_EXCELLENT,
@@ -60,8 +62,12 @@ async def test_entity_initialization(hass: HomeAssistant):
 
     for lvl, icon in levels.items():
         # pylint: disable=cell-var-from-loop
-        with patch.object(IaqukSensor, "state", new_callable=lambda: lvl):
+        with patch.object(Iaquk, "iaq_level", new_callable=lambda: lvl):
+            controller = Iaquk(hass, "test", "Test", {"": "sensor.test_monitored"})
             entity = IaqukSensor(controller, SENSOR_LEVEL)
+            entity.hass = hass
+            await entity.async_update()
+
             assert entity.state == lvl
             assert entity.icon == icon