Add state_class attribute to keba integration (#54271)

Co-authored-by: Joakim Sørensen <hi@ludeeus.dev>

homeassistant/components/keba/sensor.py

@@ -1,10 +1,18 @@
 """Support for KEBA charging station sensors."""
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import (
-from homeassistant.const import (
+    DEVICE_CLASS_CURRENT,
+    DEVICE_CLASS_ENERGY,
     DEVICE_CLASS_POWER,
+    STATE_CLASS_MEASUREMENT,
+    SensorEntity,
+    SensorEntityDescription,
+)
+from homeassistant.const import (
     ELECTRIC_CURRENT_AMPERE,
     ENERGY_KILO_WATT_HOUR,
+    POWER_KILO_WATT,
 )
+from homeassistant.util import dt
 
 from . import DOMAIN
 
@@ -19,44 +27,56 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
     sensors = [
         KebaSensor(
             keba,
-            "Curr user",
-            "Max Current",
             "max_current",
-            "mdi:flash",
+            SensorEntityDescription(
-            ELECTRIC_CURRENT_AMPERE,
+                key="Curr user",
+                name="Max Current",
+                native_unit_of_measurement=ELECTRIC_CURRENT_AMPERE,
+                device_class=DEVICE_CLASS_CURRENT,
+            ),
         ),
         KebaSensor(
             keba,
-            "Setenergy",
-            "Energy Target",
             "energy_target",
-            "mdi:gauge",
+            SensorEntityDescription(
-            ENERGY_KILO_WATT_HOUR,
+                key="Setenergy",
+                name="Energy Target",
+                native_unit_of_measurement=ENERGY_KILO_WATT_HOUR,
+                device_class=DEVICE_CLASS_ENERGY,
+            ),
         ),
         KebaSensor(
             keba,
-            "P",
-            "Charging Power",
             "charging_power",
-            "mdi:flash",
+            SensorEntityDescription(
-            "kW",
+                key="P",
-            DEVICE_CLASS_POWER,
+                name="Charging Power",
+                native_unit_of_measurement=POWER_KILO_WATT,
+                device_class=DEVICE_CLASS_POWER,
+                state_class=STATE_CLASS_MEASUREMENT,
+            ),
         ),
         KebaSensor(
             keba,
-            "E pres",
-            "Session Energy",
             "session_energy",
-            "mdi:gauge",
+            SensorEntityDescription(
-            ENERGY_KILO_WATT_HOUR,
+                key="E pres",
+                name="Session Energy",
+                native_unit_of_measurement=ENERGY_KILO_WATT_HOUR,
+                device_class=DEVICE_CLASS_ENERGY,
+            ),
         ),
         KebaSensor(
             keba,
-            "E total",
-            "Total Energy",
             "total_energy",
-            "mdi:gauge",
+            SensorEntityDescription(
-            ENERGY_KILO_WATT_HOUR,
+                key="E total",
+                name="Total Energy",
+                native_unit_of_measurement=ENERGY_KILO_WATT_HOUR,
+                device_class=DEVICE_CLASS_ENERGY,
+                state_class=STATE_CLASS_MEASUREMENT,
+                last_reset=dt.utc_from_timestamp(0),
+            ),
         ),
     ]
     async_add_entities(sensors)
@@ -65,53 +85,23 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
 class KebaSensor(SensorEntity):
     """The entity class for KEBA charging stations sensors."""
 
-    def __init__(self, keba, key, name, entity_type, icon, unit, device_class=None):
+    _attr_should_poll = False
+
+    def __init__(
+        self,
+        keba,
+        entity_type,
+        description: SensorEntityDescription,
+    ):
         """Initialize the KEBA Sensor."""
         self._keba = keba
-        self._key = key
+        self.entity_description = description
-        self._name = name
         self._entity_type = entity_type
-        self._icon = icon
-        self._unit = unit
-        self._device_class = device_class
 
-        self._state = None
+        self._attr_name = f"{keba.device_name} {description.name}"
-        self._attributes = {}
+        self._attr_unique_id = f"{keba.device_id}_{entity_type}"
 
-    @property
+        self._attributes: dict[str, str] = {}
-    def should_poll(self):
-        """Deactivate polling. Data updated by KebaHandler."""
-        return False
-
-    @property
-    def unique_id(self):
-        """Return the unique ID of the binary sensor."""
-        return f"{self._keba.device_id}_{self._entity_type}"
-
-    @property
-    def name(self):
-        """Return the name of the device."""
-        return f"{self._keba.device_name} {self._name}"
-
-    @property
-    def device_class(self):
-        """Return the class of this sensor."""
-        return self._device_class
-
-    @property
-    def icon(self):
-        """Icon to use in the frontend, if any."""
-        return self._icon
-
-    @property
-    def native_value(self):
-        """Return the state of the sensor."""
-        return self._state
-
-    @property
-    def native_unit_of_measurement(self):
-        """Get the unit of measurement."""
-        return self._unit
 
     @property
     def extra_state_attributes(self):
@@ -120,9 +110,9 @@ class KebaSensor(SensorEntity):
 
     async def async_update(self):
         """Get latest cached states from the device."""
-        self._state = self._keba.get_value(self._key)
+        self._attr_native_value = self._keba.get_value(self.entity_description.key)
 
-        if self._key == "P":
+        if self.entity_description.key == "P":
             self._attributes["power_factor"] = self._keba.get_value("PF")
             self._attributes["voltage_u1"] = str(self._keba.get_value("U1"))
             self._attributes["voltage_u2"] = str(self._keba.get_value("U2"))
@@ -130,7 +120,7 @@ class KebaSensor(SensorEntity):
             self._attributes["current_i1"] = str(self._keba.get_value("I1"))
             self._attributes["current_i2"] = str(self._keba.get_value("I2"))
             self._attributes["current_i3"] = str(self._keba.get_value("I3"))
-        elif self._key == "Curr user":
+        elif self.entity_description.key == "Curr user":
             self._attributes["max_current_hardware"] = self._keba.get_value("Curr HW")
 
     def update_callback(self):