From 2b029579e1eecb96137b77db2393099ed1541b0f Mon Sep 17 00:00:00 2001
From: Tanel Vakker <tanel@vakker.org>
Date: Sat, 21 Dec 2024 11:46:53 +0200
Subject: [PATCH] Reworks sensor to include state and device classes. Sensor
 names are now human readable and defines entity_id to have qw_ in front for
 easy grouping.

---
 custom_components/qilowatt/sensor.py | 111 +++++++++++++++++++++------
 1 file changed, 87 insertions(+), 24 deletions(-)

diff --git a/custom_components/qilowatt/sensor.py b/custom_components/qilowatt/sensor.py
index ef3d606..7e1b8a9 100644
--- a/custom_components/qilowatt/sensor.py
+++ b/custom_components/qilowatt/sensor.py
@@ -1,18 +1,64 @@
 # custom_components/qilowatt/sensor.py
 
 import logging
-
-from homeassistant.components.sensor import Entity
+from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity import DeviceInfo, Entity
+from homeassistant.helpers.entity import DeviceInfo, async_generate_entity_id
 from qilowatt import WorkModeCommand
 
 from .const import CONF_INVERTER_ID, DATA_CLIENT, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
+ENTITY_ID_FORMAT = "sensor.{}"
+
+# Define the metadata for each field
+WORKMODE_FIELDS = {
+    "Mode": {
+        "name": "Mode",
+        "unit_of_measurement": None,
+        "device_class": None,
+        "state_class": None,
+    },
+    "_source": {
+        "name": "Source",
+        "unit_of_measurement": None,
+        "device_class": None,
+        "state_class": None,
+    },
+    "BatterySoc": {
+        "name": "Battery State of Charge",
+        "unit_of_measurement": "%",
+        "device_class": "battery",
+        "state_class": "measurement",
+    },
+    "PowerLimit": {
+        "name": "Power Limit",
+        "unit_of_measurement": "W",
+        "device_class": "power",
+        "state_class": "measurement",
+    },
+    "PeakShaving": {
+        "name": "Peak Shaving",
+        "unit_of_measurement": "W",
+        "device_class": "power",
+        "state_class": "measurement",
+    },
+    "ChargeCurrent": {
+        "name": "Charge Current",
+        "unit_of_measurement": "A",
+        "device_class": "current",
+        "state_class": "measurement",
+    },
+    "DischargeCurrent": {
+        "name": "Discharge Current",
+        "unit_of_measurement": "A",
+        "device_class": "current",
+        "state_class": "measurement",
+    },
+}
 
 async def async_setup_entry(
     hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities
@@ -22,36 +68,38 @@ async def async_setup_entry(
 
     # Add sensors for WORKMODE commands
     workmode_sensors = []
-    workmode_fields = [
-        "Mode",
-        "_source",
-        "BatterySoc",
-        "PowerLimit",
-        "PeakShaving",
-        "ChargeCurrent",
-        "DischargeCurrent",
-    ]
-    for field in workmode_fields:
+    for field, metadata in WORKMODE_FIELDS.items():
+        entity_description = SensorEntityDescription(
+            key=field,
+            name=metadata["name"],
+            unit_of_measurement=metadata["unit_of_measurement"],
+            device_class=metadata["device_class"],
+            state_class=metadata["state_class"],
+        )
         sensor = WorkModeSensor(
+            hass,
             inverter_id,
-            field,
+            entity_description,
             config_entry,
         )
         workmode_sensors.append(sensor)
 
     async_add_entities(workmode_sensors, update_before_add=True)
 
-
-class WorkModeSensor(Entity):
+class WorkModeSensor(SensorEntity):
     """Sensor for WORKMODE command fields."""
 
-    def __init__(self, inverter_id, field_name, entry) -> None:
+    def __init__(self, hass: HomeAssistant, inverter_id, entity_description: SensorEntityDescription, entry) -> None:
+        self.hass = hass
         self._inverter_id = inverter_id
-        self._field_name = field_name
+        self.entity_description = entity_description
         self.entry = entry
-        self._name = field_name
-        self._unique_id = f"{inverter_id}_{field_name}"
+        self._name = entity_description.name
+        self._unique_id = f"{inverter_id}_{entity_description.key}"
         self._state = None
+        self.entity_id = async_generate_entity_id(
+            ENTITY_ID_FORMAT, f"qw_{entity_description.key}", hass.states.async_entity_ids()
+        )
 
     @property
     def name(self):
@@ -60,8 +108,8 @@ def name(self):
 
     @property
     def unique_id(self):
-        """Return the name of the sensor."""
-        return self._name
+        """Return the unique ID of the sensor."""
+        return self._unique_id
 
     @property
     def device_info(self) -> DeviceInfo:
@@ -79,6 +127,21 @@ def state(self):
         """Return the state of the sensor."""
         return self._state
 
+    @property
+    def unit_of_measurement(self):
+        """Return the unit of measurement of the sensor."""
+        return self.entity_description.unit_of_measurement
+
+    @property
+    def device_class(self):
+        """Return the device class of the sensor."""
+        return self.entity_description.device_class
+
+    @property
+    def state_class(self):
+        """Return the state class of the sensor."""
+        return self.entity_description.state_class
+
     async def async_added_to_hass(self):
         """Register dispatcher to listen for WORKMODE updates."""
         self.async_on_remove(
@@ -92,6 +155,6 @@ async def async_added_to_hass(self):
     async def _handle_workmode_update(self, command: WorkModeCommand):
         """Handle WORKMODE command updates."""
         _LOGGER.debug(f"WorkModeSensor '{self._name}' handling update.")
-        value = getattr(command, self._field_name, None)
+        value = getattr(command, self.entity_description.key, None)
         self._state = value
-        self.async_schedule_update_ha_state()
+        self.async_schedule_update_ha_state()
\ No newline at end of file