From e50ffa152fafa6e17c06b55c143a3f6c7ff3dd63 Mon Sep 17 00:00:00 2001
From: Daniel Raper <me@danr.uk>
Date: Wed, 3 Jan 2024 17:04:32 +0000
Subject: [PATCH] Give OhmeChargeSchedulesCoordinator a kick to update

---
 custom_components/ohme/const.py  |  2 +-
 custom_components/ohme/number.py | 21 +++++++++++++++++++--
 custom_components/ohme/time.py   | 24 +++++++++++++++++++++---
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/custom_components/ohme/const.py b/custom_components/ohme/const.py
index b98528b..4f03218 100644
--- a/custom_components/ohme/const.py
+++ b/custom_components/ohme/const.py
@@ -1,7 +1,7 @@
 """Component constants"""
 DOMAIN = "ohme"
 USER_AGENT = "dan-r-homeassistant-ohme"
-INTEGRATION_VERSION = "0.3.0"
+INTEGRATION_VERSION = "0.3.1"
 
 DATA_CLIENT = "client"
 DATA_COORDINATORS = "coordinators"
diff --git a/custom_components/ohme/number.py b/custom_components/ohme/number.py
index 44d80ce..cc52524 100644
--- a/custom_components/ohme/number.py
+++ b/custom_components/ohme/number.py
@@ -43,6 +43,20 @@ def __init__(self, coordinator, coordinator_schedules, hass: HomeAssistant, clie
 
         self._attr_device_info = client.get_device_info()
 
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        await super().async_added_to_hass()
+        self.async_on_remove(
+            self.coordinator.async_add_listener(
+                self._handle_coordinator_update, None
+            )
+        )
+        self.async_on_remove(
+            self.coordinator_schedules.async_add_listener(
+                self._handle_coordinator_update, None
+            )
+        )
+        
     @property
     def unique_id(self):
         """The unique ID of the switch."""
@@ -65,8 +79,8 @@ def icon(self):
         """Icon of the sensor."""
         return "mdi:battery-heart"
 
-    @property
-    def native_value(self):
+    @callback
+    def _handle_coordinator_update(self) -> None:
         """Get value from data returned from API by coordinator"""
         # Set with the same logic as reading
         if session_in_progress(self.coordinator.data):
@@ -75,4 +89,7 @@ def native_value(self):
             target = round(self.coordinator_schedules.data['targetPercent'])
 
         self._state = target if target > 0 else None
+
+    @property
+    def native_value(self):
         return self._state
diff --git a/custom_components/ohme/time.py b/custom_components/ohme/time.py
index 6eb85b4..6afc7e2 100644
--- a/custom_components/ohme/time.py
+++ b/custom_components/ohme/time.py
@@ -45,6 +45,20 @@ def __init__(self, coordinator, coordinator_schedules, hass: HomeAssistant, clie
             "number.{}", "ohme_target_time", hass=hass)
 
         self._attr_device_info = client.get_device_info()
+    
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        await super().async_added_to_hass()
+        self.async_on_remove(
+            self.coordinator.async_add_listener(
+                self._handle_coordinator_update, None
+            )
+        )
+        self.async_on_remove(
+            self.coordinator_schedules.async_add_listener(
+                self._handle_coordinator_update, None
+            )
+        )
 
     @property
     def unique_id(self):
@@ -68,10 +82,10 @@ def icon(self):
         """Icon of the sensor."""
         return "mdi:alarm-check"
 
-    @property
-    def native_value(self):
+    @callback
+    def _handle_coordinator_update(self) -> None:
         """Get value from data returned from API by coordinator"""
-        # Set with the same logic as reading
+        # Read with the same logic as setting
         target = None
         if session_in_progress(self.coordinator.data):
             target = self.coordinator.data['appliedRule']['targetTime']
@@ -84,4 +98,8 @@ def native_value(self):
                 minute=(target % 3600) // 60,
                 second=0
             )
+        self.async_write_ha_state()
+
+    @property
+    def native_value(self):
         return self._state