From 6d2cb94284b872b0b3dddbd536800a4a91fadabe Mon Sep 17 00:00:00 2001
From: tiger <superheaoz@hotmail.com>
Date: Tue, 17 Dec 2024 19:24:09 +0800
Subject: [PATCH] feat: add state_class for sensor

---
 custom_components/xiaomi_home/miot/miot_spec.py | 2 ++
 custom_components/xiaomi_home/sensor.py         | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/custom_components/xiaomi_home/miot/miot_spec.py b/custom_components/xiaomi_home/miot/miot_spec.py
index 3df70f1..195f674 100644
--- a/custom_components/xiaomi_home/miot/miot_spec.py
+++ b/custom_components/xiaomi_home/miot/miot_spec.py
@@ -79,6 +79,7 @@ class MIoTSpecBase:
     # External params
     platform: str
     device_class: any
+    state_class: any
     icon: str
     external_unit: any
 
@@ -96,6 +97,7 @@ class MIoTSpecBase:
 
         self.platform = None
         self.device_class = None
+        self.state_class = None
         self.icon = None
         self.external_unit = None
 
diff --git a/custom_components/xiaomi_home/sensor.py b/custom_components/xiaomi_home/sensor.py
index 7d2e074..3ef80b5 100644
--- a/custom_components/xiaomi_home/sensor.py
+++ b/custom_components/xiaomi_home/sensor.py
@@ -105,6 +105,9 @@ class Sensor(MIoTPropertyEntity, SensorEntity):
         # Set icon
         if spec.icon:
             self._attr_icon = spec.icon
+        # Set state_class
+        if spec.state_class:
+            self._attr_state_class = spec.state_class
 
     @property
     def native_value(self) -> any: