Fix rain count sensors' state class of Ecowitt

homeassistant/components/ecowitt/sensor.py

@@ -4,6 +4,7 @@
 
 import dataclasses
 from datetime import datetime
+import logging
 from typing import Final
 
 from aioecowitt import EcoWittSensor, EcoWittSensorTypes
@@ -39,6 +40,9 @@
 from . import EcowittConfigEntry
 from .entity import EcowittEntity
 
+_LOGGER = logging.getLogger(__name__)
+
+
 _METRIC: Final = (
     EcoWittSensorTypes.TEMPERATURE_C,
     EcoWittSensorTypes.RAIN_COUNT_MM,
@@ -57,6 +61,40 @@
 )
 
 
+_RAIN_COUNT_SENSORS_STATE_CLASS_MAPPING: Final = {
+    "eventrainin": SensorStateClass.TOTAL_INCREASING,
+    "hourlyrainin": None,
+    "totalrainin": SensorStateClass.TOTAL_INCREASING,
+    "dailyrainin": SensorStateClass.TOTAL_INCREASING,
+    "weeklyrainin": SensorStateClass.TOTAL_INCREASING,
+    "monthlyrainin": SensorStateClass.TOTAL_INCREASING,
+    "yearlyrainin": SensorStateClass.TOTAL_INCREASING,
+    "last24hrainin": None,
+    "eventrainmm": SensorStateClass.TOTAL_INCREASING,
+    "hourlyrainmm": None,
+    "totalrainmm": SensorStateClass.TOTAL_INCREASING,
+    "dailyrainmm": SensorStateClass.TOTAL_INCREASING,
+    "weeklyrainmm": SensorStateClass.TOTAL_INCREASING,
+    "monthlyrainmm": SensorStateClass.TOTAL_INCREASING,
+    "yearlyrainmm": SensorStateClass.TOTAL_INCREASING,
+    "last24hrainmm": None,
+    "erain_piezo": SensorStateClass.TOTAL_INCREASING,
+    "hrain_piezo": None,
+    "drain_piezo": SensorStateClass.TOTAL_INCREASING,
+    "wrain_piezo": SensorStateClass.TOTAL_INCREASING,
+    "mrain_piezo": SensorStateClass.TOTAL_INCREASING,
+    "yrain_piezo": SensorStateClass.TOTAL_INCREASING,
+    "last24hrain_piezo": None,
+    "erain_piezomm": SensorStateClass.TOTAL_INCREASING,
+    "hrain_piezomm": None,
+    "drain_piezomm": SensorStateClass.TOTAL_INCREASING,
+    "wrain_piezomm": SensorStateClass.TOTAL_INCREASING,
+    "mrain_piezomm": SensorStateClass.TOTAL_INCREASING,
+    "yrain_piezomm": SensorStateClass.TOTAL_INCREASING,
+    "last24hrain_piezomm": None,
+}
+
+
 ECOWITT_SENSORS_MAPPING: Final = {
     EcoWittSensorTypes.HUMIDITY: SensorEntityDescription(
         key="HUMIDITY",
@@ -285,15 +323,15 @@ def _new_sensor(sensor: EcoWittSensor) -> None:
             name=sensor.name,
         )
 
-        # Only total rain needs state class for long-term statistics
-        if sensor.key in (
-            "totalrainin",
-            "totalrainmm",
+        if sensor.stype in (
+            EcoWittSensorTypes.RAIN_COUNT_INCHES,
+            EcoWittSensorTypes.RAIN_COUNT_MM,
         ):
-            description = dataclasses.replace(
-                description,
-                state_class=SensorStateClass.TOTAL_INCREASING,
-            )
+            if sensor.key not in _RAIN_COUNT_SENSORS_STATE_CLASS_MAPPING:
+                _LOGGER.warning("Unknown rain count sensor: %s", sensor.key)
+                return
+            state_class = _RAIN_COUNT_SENSORS_STATE_CLASS_MAPPING[sensor.key]
+            description = dataclasses.replace(description, state_class=state_class)
 
         async_add_entities([EcowittSensorEntity(sensor, description)])