From 9f1c02b179904b6b9816807fe5c17039b20abf2a Mon Sep 17 00:00:00 2001
From: "Shaun J.V.Nieuwenhuizen" <shaunjvn@gmail.com>
Date: Sat, 4 May 2024 12:50:50 +0200
Subject: [PATCH] possible fix for #57

---
 custom_components/solis_modbus/const.py      |  2 +-
 custom_components/solis_modbus/manifest.json |  2 +-
 custom_components/solis_modbus/number.py     | 10 ++++++---
 custom_components/solis_modbus/sensor.py     | 22 +++++++++++++-------
 4 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/custom_components/solis_modbus/const.py b/custom_components/solis_modbus/const.py
index 5334c4d..bea6e34 100644
--- a/custom_components/solis_modbus/const.py
+++ b/custom_components/solis_modbus/const.py
@@ -1,6 +1,6 @@
 DOMAIN = "solis_modbus"
 CONTROLLER = "modbus_controller"
-VERSION = "1.4.8"
+VERSION = "1.4.9"
 POLL_INTERVAL_SECONDS = 15
 MANUFACTURER = "Solis"
 MODEL = "S6"
diff --git a/custom_components/solis_modbus/manifest.json b/custom_components/solis_modbus/manifest.json
index 9dd9e0d..8e331a6 100644
--- a/custom_components/solis_modbus/manifest.json
+++ b/custom_components/solis_modbus/manifest.json
@@ -10,5 +10,5 @@
   "issue_tracker": "https://github.com/Pho3niX90/solis_modbus/issues",
   "quality_scale": "silver",
   "requirements": ["pymodbus>=3.6.8"],
-  "version": "1.4.8"
+  "version": "1.4.9"
 }
diff --git a/custom_components/solis_modbus/number.py b/custom_components/solis_modbus/number.py
index 0305bcd..560929c 100644
--- a/custom_components/solis_modbus/number.py
+++ b/custom_components/solis_modbus/number.py
@@ -11,7 +11,7 @@
 from typing import List
 
 from homeassistant.components.number import NumberEntity, NumberMode
-from homeassistant.components.sensor import SensorDeviceClass
+from homeassistant.components.sensor import SensorDeviceClass, RestoreSensor
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     UnitOfElectricCurrent, PERCENTAGE, UnitOfPower, )
@@ -102,7 +102,7 @@ async def get_modbus_updates(hass, controller):
     # fmt: on
 
 
-class SolisNumberEntity(NumberEntity):
+class SolisNumberEntity(RestoreSensor, NumberEntity):
     """Representation of a Number entity."""
 
     def __init__(self, hass, modbus_controller, entity_definition):
@@ -121,6 +121,7 @@ def __init__(self, hass, modbus_controller, entity_definition):
         self._attr_native_value = entity_definition.get("default", None)
         self._attr_assumed_state = entity_definition.get("assumed", False)
         self._attr_available = False
+        self.is_added_to_hass = False
         self._attr_device_class = entity_definition.get("device_class", None)
         self._attr_icon = entity_definition.get("icon", None)
         self._attr_mode = entity_definition.get("mode", NumberMode.AUTO)
@@ -133,7 +134,10 @@ def __init__(self, hass, modbus_controller, entity_definition):
 
     async def async_added_to_hass(self) -> None:
         await super().async_added_to_hass()
-        _LOGGER.info(f"async_added_to_hass {self._attr_name},  {self.entity_id},  {self.unique_id}")
+        state = await self.async_get_last_sensor_data()
+        if state:
+            self._attr_native_value = state.native_value
+        self.is_added_to_hass = True
 
     def update(self):
         """Update Modbus data periodically."""
diff --git a/custom_components/solis_modbus/sensor.py b/custom_components/solis_modbus/sensor.py
index 67bd1b5..3815089 100644
--- a/custom_components/solis_modbus/sensor.py
+++ b/custom_components/solis_modbus/sensor.py
@@ -840,10 +840,11 @@ def update(self):
                 else:
                     n_value = 0
 
-            self._attr_available = True
-            self._attr_native_value = n_value * self._display_multiplier
-            self._state = n_value * self._display_multiplier
-            self.schedule_update_ha_state()
+            if n_value is not None:
+                self._attr_available = True
+                self._attr_native_value = n_value * self._display_multiplier
+                self._state = n_value * self._display_multiplier
+                self.schedule_update_ha_state()
 
         except ValueError as e:
             _LOGGER.error(e)
@@ -902,10 +903,13 @@ async def async_added_to_hass(self) -> None:
 
     def update(self):
         """Update the sensor value."""
+
         try:
             if not self.is_added_to_hass:
                 return
 
+            n_value = None
+
             if '33027' in self._register:
                 hours = self._hass.data[DOMAIN]['values'][str(int(self._register[0]) - 2)]
                 minutes = self._hass.data[DOMAIN]['values'][str(int(self._register[0]) - 1)]
@@ -917,10 +921,12 @@ def update(self):
             else:
                 n_value = get_value(self)
 
-            self._attr_available = True
-            self._attr_native_value = n_value * self._display_multiplier
-            self._state = n_value * self._display_multiplier
-            self.schedule_update_ha_state()
+            if n_value is not None:
+                self._attr_available = True
+                self._attr_native_value = n_value * self._display_multiplier
+                self._state = n_value * self._display_multiplier
+                self.schedule_update_ha_state()
+
         except ValueError as e:
             _LOGGER.error(e)
             # Handle communication or reading errors