feat: potentially fixed absent state class on sensors in Home Assista…

…nt 2022.3.0+
alryaz · Apr 7, 2022 · 7bfe687 · 7bfe687
1 parent f6c1a31
commit 7bfe687
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 28 deletions.
diff --git a/custom_components/hekr/base_platform.py b/custom_components/hekr/base_platform.py
@@ -13,7 +13,6 @@
 import voluptuous as vol
 from homeassistant import config_entries
 from homeassistant.components.binary_sensor.device_condition import DEVICE_CLASS_NONE
-from homeassistant.components.sensor import ATTR_STATE_CLASS
 from homeassistant.const import (
     STATE_UNKNOWN,
     STATE_OK,
@@ -264,11 +263,6 @@ def device_state_attributes(self) -> Optional[dict]:
         if attributes:
             base_attributes.update(attributes)
 
-        if ATTR_STATE_CLASS not in base_attributes:
-            state_class = self.state_class
-            if state_class is not None:
-                base_attributes[ATTR_STATE_CLASS] = state_class
-
         return base_attributes
 
     @property
@@ -290,10 +284,6 @@ def command_receive(self) -> str:
     def device_class(self) -> Optional[str]:
         return self._config.get(ATTR_DEVICE_CLASS, DEVICE_CLASS_NONE)
 
-    @property
-    def state_class(self) -> Optional[str]:
-        return self._config.get(ATTR_STATE_CLASS)
-
     @property
     def device_info(self) -> Optional[Dict[str, Any]]:
         return self._data.get_device_info_dict(self._device_id)

diff --git a/custom_components/hekr/sensor.py b/custom_components/hekr/sensor.py
@@ -10,19 +10,28 @@
 import logging
 from typing import Optional
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, DOMAIN as PLATFORM_DOMAIN
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    DOMAIN as PLATFORM_DOMAIN,
+    SensorEntity,
+    ATTR_STATE_CLASS,
+)
 from homeassistant.helpers.restore_state import RestoreEntity
 
 from .base_platform import HekrEntity, create_platform_basics
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class HekrSensor(HekrEntity, RestoreEntity):
+class HekrSensor(HekrEntity, RestoreEntity, SensorEntity):
     @property
     def unique_id(self) -> Optional[str]:
         return "_".join((self._device_id, PLATFORM_DOMAIN, self._ent_type))
 
+    @property
+    def state_class(self) -> Optional[str]:
+        return self._config.get(ATTR_STATE_CLASS)
+
 
 PLATFORM_SCHEMA, async_setup_platform, async_setup_entry = create_platform_basics(
     logger=_LOGGER,

diff --git a/custom_components/hekr/supported_protocols.py b/custom_components/hekr/supported_protocols.py
@@ -5,10 +5,16 @@
     "POWER_METER",
 ]
 
+from hekrapi.protocols.power_meter import (
+    CurrentWarning,
+    PROTOCOL as PROTOCOL_POWER_METER,
+    PowerSupplyWarning,
+    VoltageWarning,
+)
+from hekrapi.protocols.power_socket import PROTOCOL as PROTOCOL_POWER_SOCKET
 from homeassistant.components.sensor import (
     ATTR_STATE_CLASS,
-    STATE_CLASS_MEASUREMENT,
-    STATE_CLASS_TOTAL_INCREASING,
+    SensorStateClass,
 )
 from homeassistant.components.switch import DEVICE_CLASS_SWITCH
 from homeassistant.const import (
@@ -33,13 +39,6 @@
     STATE_PROBLEM,
 )
 
-from hekrapi.protocols.power_meter import (
-    CurrentWarning,
-    PROTOCOL as PROTOCOL_POWER_METER,
-    PowerSupplyWarning,
-    VoltageWarning,
-)
-from hekrapi.protocols.power_socket import PROTOCOL as PROTOCOL_POWER_SOCKET
 from .const import (
     ATTR_MONITORED,
     PROTOCOL_CMD_RECEIVE,
@@ -170,7 +169,7 @@ def power_meter_attribute_filter(attributes: dict) -> dict:
             ATTR_STATE: "current_energy_consumption",
             ATTR_UNIT_OF_MEASUREMENT: POWER_WATT,
             ATTR_DEVICE_CLASS: DEVICE_CLASS_POWER,
-            ATTR_STATE_CLASS: STATE_CLASS_MEASUREMENT,
+            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
             PROTOCOL_CMD_UPDATE: "queryDev",
             PROTOCOL_CMD_RECEIVE: "reportDev",
             PROTOCOL_DEFAULT: True,
@@ -181,7 +180,7 @@ def power_meter_attribute_filter(attributes: dict) -> dict:
             ATTR_STATE: "total_energy_consumed",
             ATTR_UNIT_OF_MEASUREMENT: ENERGY_KILO_WATT_HOUR,
             ATTR_DEVICE_CLASS: DEVICE_CLASS_ENERGY,
-            ATTR_STATE_CLASS: STATE_CLASS_TOTAL_INCREASING,
+            ATTR_STATE_CLASS: SensorStateClass.TOTAL_INCREASING,
             PROTOCOL_CMD_UPDATE: "queryDev",
             PROTOCOL_CMD_RECEIVE: "reportDev",
             PROTOCOL_DEFAULT: True,
@@ -193,7 +192,7 @@ def power_meter_attribute_filter(attributes: dict) -> dict:
             ATTR_STATE: "mean_voltage",
             ATTR_UNIT_OF_MEASUREMENT: ELECTRIC_POTENTIAL_VOLT,
             ATTR_DEVICE_CLASS: DEVICE_CLASS_VOLTAGE,
-            ATTR_STATE_CLASS: STATE_CLASS_MEASUREMENT,
+            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
             ATTR_MONITORED: [
                 "voltage_1",
                 "voltage_2",
@@ -210,7 +209,7 @@ def power_meter_attribute_filter(attributes: dict) -> dict:
             ATTR_STATE: "total_current",
             ATTR_UNIT_OF_MEASUREMENT: ELECTRIC_CURRENT_AMPERE,
             ATTR_DEVICE_CLASS: DEVICE_CLASS_CURRENT,
-            ATTR_STATE_CLASS: STATE_CLASS_MEASUREMENT,
+            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
             ATTR_MONITORED: [
                 "mean_current",
                 "current_1",
@@ -228,7 +227,7 @@ def power_meter_attribute_filter(attributes: dict) -> dict:
             ATTR_STATE: "total_power_factor",
             ATTR_UNIT_OF_MEASUREMENT: PERCENTAGE,
             ATTR_DEVICE_CLASS: DEVICE_CLASS_POWER_FACTOR,
-            ATTR_STATE_CLASS: STATE_CLASS_MEASUREMENT,
+            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
             ATTR_MONITORED: ["power_factor_1", "power_factor_2", "power_factor_3"],
             PROTOCOL_CMD_UPDATE: "queryData",
             PROTOCOL_CMD_RECEIVE: "reportData",
@@ -240,7 +239,7 @@ def power_meter_attribute_filter(attributes: dict) -> dict:
             ATTR_STATE: "total_active_power",
             ATTR_UNIT_OF_MEASUREMENT: POWER_WATT,
             ATTR_DEVICE_CLASS: DEVICE_CLASS_POWER,
-            ATTR_STATE_CLASS: STATE_CLASS_MEASUREMENT,
+            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
             ATTR_MONITORED: ["active_power_1", "active_power_2", "active_power_3"],
             PROTOCOL_CMD_UPDATE: "queryData",
             PROTOCOL_CMD_RECEIVE: "reportData",
@@ -251,7 +250,7 @@ def power_meter_attribute_filter(attributes: dict) -> dict:
             ATTR_ICON: "mdi:flash-outline",
             ATTR_STATE: "total_reactive_power",
             ATTR_UNIT_OF_MEASUREMENT: "kVar",
-            ATTR_STATE_CLASS: STATE_CLASS_MEASUREMENT,
+            ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
             ATTR_MONITORED: [
                 "reactive_power_1",
                 "reactive_power_2",