From 6d7ad8903f74959c6e3309907c77bc256bff5555 Mon Sep 17 00:00:00 2001
From: Petru Paler <petru@paler.net>
Date: Tue, 17 Aug 2021 09:19:44 +0100
Subject: [PATCH] Energy support for Solax inverters (#54654)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>
---
 homeassistant/components/solax/sensor.py | 51 ++++++++++++++++++++++--
 1 file changed, 47 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/solax/sensor.py b/homeassistant/components/solax/sensor.py
index 4d1652e8b12..7854142c32b 100644
--- a/homeassistant/components/solax/sensor.py
+++ b/homeassistant/components/solax/sensor.py
@@ -6,8 +6,23 @@ from solax import real_time_api
 from solax.inverter import InverterError
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
-from homeassistant.const import CONF_IP_ADDRESS, CONF_PORT, TEMP_CELSIUS
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    STATE_CLASS_MEASUREMENT,
+    STATE_CLASS_TOTAL_INCREASING,
+    SensorEntity,
+)
+from homeassistant.const import (
+    CONF_IP_ADDRESS,
+    CONF_PORT,
+    DEVICE_CLASS_BATTERY,
+    DEVICE_CLASS_CURRENT,
+    DEVICE_CLASS_ENERGY,
+    DEVICE_CLASS_POWER,
+    DEVICE_CLASS_TEMPERATURE,
+    DEVICE_CLASS_VOLTAGE,
+    TEMP_CELSIUS,
+)
 from homeassistant.exceptions import PlatformNotReady
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.event import async_track_time_interval
@@ -34,10 +49,28 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
     async_track_time_interval(hass, endpoint.async_refresh, SCAN_INTERVAL)
     devices = []
     for sensor, (idx, unit) in api.inverter.sensor_map().items():
+        device_class = state_class = None
         if unit == "C":
+            device_class = DEVICE_CLASS_TEMPERATURE
+            state_class = STATE_CLASS_MEASUREMENT
             unit = TEMP_CELSIUS
+        elif unit == "kWh":
+            device_class = DEVICE_CLASS_ENERGY
+            state_class = STATE_CLASS_TOTAL_INCREASING
+        elif unit == "V":
+            device_class = DEVICE_CLASS_VOLTAGE
+            state_class = STATE_CLASS_MEASUREMENT
+        elif unit == "A":
+            device_class = DEVICE_CLASS_CURRENT
+            state_class = STATE_CLASS_MEASUREMENT
+        elif unit == "W":
+            device_class = DEVICE_CLASS_POWER
+            state_class = STATE_CLASS_MEASUREMENT
+        elif unit == "%":
+            device_class = DEVICE_CLASS_BATTERY
+            state_class = STATE_CLASS_MEASUREMENT
         uid = f"{serial}-{idx}"
-        devices.append(Inverter(uid, serial, sensor, unit))
+        devices.append(Inverter(uid, serial, sensor, unit, state_class, device_class))
     endpoint.sensors = devices
     async_add_entities(devices)
 
@@ -75,13 +108,23 @@ class RealTimeDataEndpoint:
 class Inverter(SensorEntity):
     """Class for a sensor."""
 
-    def __init__(self, uid, serial, key, unit):
+    def __init__(
+        self,
+        uid,
+        serial,
+        key,
+        unit,
+        state_class=None,
+        device_class=None,
+    ):
         """Initialize an inverter sensor."""
         self.uid = uid
         self.serial = serial
         self.key = key
         self.value = None
         self.unit = unit
+        self._attr_state_class = state_class
+        self._attr_device_class = device_class
 
     @property
     def native_value(self):