Fix missing state_class for quirk v2 sensors (#341)

* Fix missing state_class for quirk v2 sensors

* Update test with invalid state class

* Assert log message

Author: abmantis

tests/test_sensor.py

@@ -36,7 +36,7 @@
 from zha.application.gateway import Gateway
 from zha.application.platforms import PlatformEntity, sensor
 from zha.application.platforms.sensor import DanfossSoftwareErrorCode, UnitOfMass
-from zha.application.platforms.sensor.const import SensorDeviceClass
+from zha.application.platforms.sensor.const import SensorDeviceClass, SensorStateClass
 from zha.units import PERCENTAGE, UnitOfEnergy, UnitOfPressure, UnitOfVolume
 from zha.zigbee.device import Device
 
@@ -1245,6 +1245,34 @@ def __init__(self, *args, **kwargs) -> None:
         translation_key="last_feeding_size",
         fallback_name="Last feeding size",
     )
+    .sensor(
+        "power",
+        OppleCluster.cluster_id,
+        state_class=SensorStateClass.MEASUREMENT,
+        device_class=SensorDeviceClass.POWER,
+        fallback_name="Measurement",
+    )
+    .sensor(
+        "energy",
+        OppleCluster.cluster_id,
+        state_class=SensorStateClass.TOTAL,
+        device_class=SensorDeviceClass.ENERGY,
+        fallback_name="Measurement",
+    )
+    .sensor(
+        "energy_delivered",
+        OppleCluster.cluster_id,
+        state_class=SensorStateClass.TOTAL_INCREASING,
+        device_class=SensorDeviceClass.ENERGY,
+        fallback_name="Measurement",
+    )
+    .sensor(
+        "energy_invalid_state_class",
+        OppleCluster.cluster_id,
+        state_class=SensorDeviceClass.ENERGY,
+        device_class=SensorDeviceClass.ENERGY,
+        fallback_name="Measurement",
+    )
     .add_to_registry()
 )
 
@@ -1292,6 +1320,29 @@ async def test_last_feeding_size_sensor_v2(zha_gateway: Gateway) -> None:
     assert_state(entity, 5.0, "g")
 
 
+async def test_state_class(
+    zha_gateway: Gateway, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test quirks defined sensor."""
+
+    zha_device, cluster = await zigpy_device_aqara_sensor_v2_mock(zha_gateway)
+    assert isinstance(zha_device.device, CustomDeviceV2)
+    power_entity = get_entity(zha_device, platform=Platform.SENSOR, qualifier="power")
+    energy_entity = get_entity(zha_device, platform=Platform.SENSOR, qualifier="energy")
+    energy_delivered_entity = get_entity(
+        zha_device, platform=Platform.SENSOR, qualifier="energy_delivered"
+    )
+    energy_invalid_state_class = get_entity(
+        zha_device, platform=Platform.SENSOR, qualifier="energy_invalid_state_class"
+    )
+
+    assert power_entity.state_class == SensorStateClass.MEASUREMENT
+    assert energy_entity.state_class == SensorStateClass.TOTAL
+    assert energy_delivered_entity.state_class == SensorStateClass.TOTAL_INCREASING
+    assert energy_invalid_state_class.state_class is None
+    assert "Quirks provided an invalid state class: energy" in caplog.text
+
+
 async def test_device_counter_sensors(zha_gateway: Gateway) -> None:
     """Test coordinator counter sensor."""
 

zha/application/platforms/sensor/__init__.py

@@ -201,6 +201,21 @@ def __init__(
             self.handle_cluster_handler_attribute_updated,
         )
 
+    def _validate_state_class(
+        self,
+        state_class_value: SensorStateClass,
+    ) -> SensorStateClass | None:
+        """Validate and return a state class."""
+        try:
+            return SensorStateClass(state_class_value.value)
+        except ValueError as ex:
+            _LOGGER.warning(
+                "Quirks provided an invalid state class: %s: %s",
+                state_class_value,
+                ex,
+            )
+            return None
+
     def _init_from_quirks_metadata(self, entity_metadata: ZCLSensorMetadata) -> None:
         """Init this entity from the quirks metadata."""
         super()._init_from_quirks_metadata(entity_metadata)
@@ -216,6 +231,10 @@ def _init_from_quirks_metadata(self, entity_metadata: ZCLSensorMetadata) -> None
                 Platform.SENSOR.value,
                 _LOGGER,
             )
+        if entity_metadata.state_class is not None:
+            self._attr_state_class = self._validate_state_class(
+                entity_metadata.state_class
+            )
         if entity_metadata.unit is not None:
             self._attr_native_unit_of_measurement = validate_unit(
                 entity_metadata.unit