fix: set state_class

domenicoblanco · Apr 2, 2024 · 3e11782 · 3e11782
1 parent 2b8a327
commit 3e11782
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/custom_components/dab_live!/entity.py b/custom_components/dab_live!/entity.py
@@ -15,22 +15,21 @@ def unique_id(self):
     def device_info(self):
         product_type = self.zone['pumps'][self.pump_id].get('ProductType')
         return {
-            "hw_version": product_type,
-            "identifiers": {(DOMAIN, self.zone['id'])},
-            "name": self.zone['name'],
-            "manufacturer": self.zone['company'],
-            "model":  MODELS.get(product_type, product_type),
-            "serial_number": self.zone['pumps'][self.pump_id].get('ProductSerialNumber'),
+            'hw_version': product_type,
+            'identifiers': {(DOMAIN, self.zone['id'])},
+            'name': self.zone['name'],
+            'manufacturer': self.zone['company'],
+            'model':  MODELS.get(product_type, product_type),
+            'serial_number': self.zone['pumps'][self.pump_id].get('ProductSerialNumber'),
             'sw_version': self.zone['pumps'][self.pump_id].get('LvVersion')
         }
 
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
         return {
-            "attribution": ATTRIBUTION,
-            "id": str(self.zone.get("id")),
-            "integration": DOMAIN,
-            "manufacturer": self.zone['company'],
-            'unit_of_measurement': self.__unit_of_measurement
+            'attribution': ATTRIBUTION,
+            'id': str(self.zone.get('id')),
+            'integration': DOMAIN,
+            'manufacturer': self.zone['company']
         }
diff --git a/custom_components/dab_live!/sensor.py b/custom_components/dab_live!/sensor.py
@@ -1,7 +1,7 @@
 from . import const
 from .entity import DABLiveEntity
 from homeassistant.const import UnitOfEnergy, UnitOfVolume
-from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
+from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass, SensorEntity
 
 
 async def async_setup_entry(hass, entry, async_add_devices):
@@ -16,7 +16,7 @@ async def async_setup_entry(hass, entry, async_add_devices):
 
     async_add_devices(entities)
 
-class DABLiveSensor(DABLiveEntity):
+class DABLiveSensor(DABLiveEntity, SensorEntity):
     """Sensor class."""
     def __init__(self, coordinator, config_entry, zone: dict, pump_id: str, sensor_type: str):
         self.zone = zone