Fix missing device_class and state_class on compensation entities (home-assistant#146115)

Co-authored-by: Robert Resch <[email protected]>

Author: Petro31

homeassistant/components/compensation/__init__.py

@@ -6,11 +6,18 @@
 import numpy as np
 import voluptuous as vol
 
-from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
+from homeassistant.components.sensor import (
+    CONF_STATE_CLASS,
+    DEVICE_CLASSES_SCHEMA as SENSOR_DEVICE_CLASSES_SCHEMA,
+    DOMAIN as SENSOR_DOMAIN,
+    STATE_CLASSES_SCHEMA as SENSOR_STATE_CLASSES_SCHEMA,
+)
 from homeassistant.const import (
     CONF_ATTRIBUTE,
+    CONF_DEVICE_CLASS,
     CONF_MAXIMUM,
     CONF_MINIMUM,
+    CONF_NAME,
     CONF_SOURCE,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
@@ -50,20 +57,23 @@ def datapoints_greater_than_degree(value: dict) -> dict:
 
 COMPENSATION_SCHEMA = vol.Schema(
     {
-        vol.Required(CONF_SOURCE): cv.entity_id,
+        vol.Optional(CONF_ATTRIBUTE): cv.string,
         vol.Required(CONF_DATAPOINTS): [
             vol.ExactSequence([vol.Coerce(float), vol.Coerce(float)])
         ],
-        vol.Optional(CONF_UNIQUE_ID): cv.string,
-        vol.Optional(CONF_ATTRIBUTE): cv.string,
-        vol.Optional(CONF_UPPER_LIMIT, default=False): cv.boolean,
-        vol.Optional(CONF_LOWER_LIMIT, default=False): cv.boolean,
-        vol.Optional(CONF_PRECISION, default=DEFAULT_PRECISION): cv.positive_int,
         vol.Optional(CONF_DEGREE, default=DEFAULT_DEGREE): vol.All(
             vol.Coerce(int),
             vol.Range(min=1, max=7),
         ),
+        vol.Optional(CONF_DEVICE_CLASS): SENSOR_DEVICE_CLASSES_SCHEMA,
+        vol.Optional(CONF_LOWER_LIMIT, default=False): cv.boolean,
+        vol.Optional(CONF_NAME): cv.string,
+        vol.Optional(CONF_PRECISION, default=DEFAULT_PRECISION): cv.positive_int,
+        vol.Required(CONF_SOURCE): cv.entity_id,
+        vol.Optional(CONF_STATE_CLASS): SENSOR_STATE_CLASSES_SCHEMA,
+        vol.Optional(CONF_UNIQUE_ID): cv.string,
         vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
+        vol.Optional(CONF_UPPER_LIMIT, default=False): cv.boolean,
     }
 )
 

homeassistant/components/compensation/sensor.py

@@ -7,15 +7,23 @@
 
 import numpy as np
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import (
+    ATTR_STATE_CLASS,
+    CONF_STATE_CLASS,
+    SensorEntity,
+)
 from homeassistant.const import (
+    ATTR_DEVICE_CLASS,
     ATTR_UNIT_OF_MEASUREMENT,
     CONF_ATTRIBUTE,
+    CONF_DEVICE_CLASS,
     CONF_MAXIMUM,
     CONF_MINIMUM,
+    CONF_NAME,
     CONF_SOURCE,
     CONF_UNIQUE_ID,
     CONF_UNIT_OF_MEASUREMENT,
+    STATE_UNAVAILABLE,
     STATE_UNKNOWN,
 )
 from homeassistant.core import (
@@ -59,24 +67,13 @@ async def async_setup_platform(
 
     source: str = conf[CONF_SOURCE]
     attribute: str | None = conf.get(CONF_ATTRIBUTE)
-    name = f"{DEFAULT_NAME} {source}"
-    if attribute is not None:
-        name = f"{name} {attribute}"
+    if not (name := conf.get(CONF_NAME)):
+        name = f"{DEFAULT_NAME} {source}"
+        if attribute is not None:
+            name = f"{name} {attribute}"
 
     async_add_entities(
-        [
-            CompensationSensor(
-                conf.get(CONF_UNIQUE_ID),
-                name,
-                source,
-                attribute,
-                conf[CONF_PRECISION],
-                conf[CONF_POLYNOMIAL],
-                conf.get(CONF_UNIT_OF_MEASUREMENT),
-                conf[CONF_MINIMUM],
-                conf[CONF_MAXIMUM],
-            )
-        ]
+        [CompensationSensor(conf.get(CONF_UNIQUE_ID), name, source, attribute, conf)]
     )
 
 
@@ -91,23 +88,27 @@ def __init__(
         name: str,
         source: str,
         attribute: str | None,
-        precision: int,
-        polynomial: np.poly1d,
-        unit_of_measurement: str | None,
-        minimum: tuple[float, float] | None,
-        maximum: tuple[float, float] | None,
+        config: dict[str, Any],
     ) -> None:
         """Initialize the Compensation sensor."""
+
+        self._attr_name = name
         self._source_entity_id = source
-        self._precision = precision
         self._source_attribute = attribute
-        self._attr_native_unit_of_measurement = unit_of_measurement
+
+        self._precision = config[CONF_PRECISION]
+        self._attr_native_unit_of_measurement = config.get(CONF_UNIT_OF_MEASUREMENT)
+
+        polynomial: np.poly1d = config[CONF_POLYNOMIAL]
         self._poly = polynomial
         self._coefficients = polynomial.coefficients.tolist()
+
         self._attr_unique_id = unique_id
-        self._attr_name = name
-        self._minimum = minimum
-        self._maximum = maximum
+        self._minimum = config[CONF_MINIMUM]
+        self._maximum = config[CONF_MAXIMUM]
+
+        self._attr_device_class = config.get(CONF_DEVICE_CLASS)
+        self._attr_state_class = config.get(CONF_STATE_CLASS)
 
     async def async_added_to_hass(self) -> None:
         """Handle added to Hass."""
@@ -137,13 +138,40 @@ def _async_compensation_sensor_state_listener(
         """Handle sensor state changes."""
         new_state: State | None
         if (new_state := event.data["new_state"]) is None:
+            _LOGGER.warning(
+                "While updating compensation %s, the new_state is None", self.name
+            )
+            self._attr_native_value = None
+            self.async_write_ha_state()
             return
 
+        if new_state.state == STATE_UNKNOWN:
+            self._attr_native_value = None
+            self.async_write_ha_state()
+            return
+
+        if new_state.state == STATE_UNAVAILABLE:
+            self._attr_available = False
+            self.async_write_ha_state()
+            return
+
+        self._attr_available = True
+
         if self.native_unit_of_measurement is None and self._source_attribute is None:
             self._attr_native_unit_of_measurement = new_state.attributes.get(
                 ATTR_UNIT_OF_MEASUREMENT
             )
 
+        if self._attr_device_class is None and (
+            device_class := new_state.attributes.get(ATTR_DEVICE_CLASS)
+        ):
+            self._attr_device_class = device_class
+
+        if self._attr_state_class is None and (
+            state_class := new_state.attributes.get(ATTR_STATE_CLASS)
+        ):
+            self._attr_state_class = state_class
+
         if self._source_attribute:
             value = new_state.attributes.get(self._source_attribute)
         else: