From 519c1fa2dad3deef8793d35dcc4da1c4554ab049 Mon Sep 17 00:00:00 2001
From: Quentame <polletquentin74@me.com>
Date: Sat, 11 Jan 2020 23:12:55 +0100
Subject: [PATCH] Update iCloud sensors when service finish its update (#30680)

* Update iCloud sensors when needed

* Add sensor should_poll
---
 homeassistant/components/icloud/__init__.py   |  4 ++--
 homeassistant/components/icloud/const.py      |  2 +-
 .../components/icloud/device_tracker.py       | 14 +++++++-------
 homeassistant/components/icloud/sensor.py     | 19 ++++++++++++++++++-
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/icloud/__init__.py b/homeassistant/components/icloud/__init__.py
index e983f5fac22..d4074db021e 100644
--- a/homeassistant/components/icloud/__init__.py
+++ b/homeassistant/components/icloud/__init__.py
@@ -46,9 +46,9 @@ from .const import (
     DEVICE_STATUS_SET,
     DOMAIN,
     ICLOUD_COMPONENTS,
+    SERVICE_UPDATE,
     STORAGE_KEY,
     STORAGE_VERSION,
-    TRACKER_UPDATE,
 )
 
 ATTRIBUTION = "Data provided by Apple iCloud"
@@ -336,7 +336,7 @@ class IcloudAccount:
                 self._devices[device_id] = IcloudDevice(self, device, status)
                 self._devices[device_id].update(status)
 
-        dispatcher_send(self.hass, TRACKER_UPDATE)
+        dispatcher_send(self.hass, SERVICE_UPDATE)
         self._fetch_interval = self._determine_interval()
         track_point_in_utc_time(
             self.hass,
diff --git a/homeassistant/components/icloud/const.py b/homeassistant/components/icloud/const.py
index ed2fc78fe6d..c2545d911df 100644
--- a/homeassistant/components/icloud/const.py
+++ b/homeassistant/components/icloud/const.py
@@ -1,7 +1,7 @@
 """iCloud component constants."""
 
 DOMAIN = "icloud"
-TRACKER_UPDATE = f"{DOMAIN}_tracker_update"
+SERVICE_UPDATE = f"{DOMAIN}_update"
 
 CONF_ACCOUNT_NAME = "account_name"
 CONF_MAX_INTERVAL = "max_interval"
diff --git a/homeassistant/components/icloud/device_tracker.py b/homeassistant/components/icloud/device_tracker.py
index 511ce7f9447..79627eec4aa 100644
--- a/homeassistant/components/icloud/device_tracker.py
+++ b/homeassistant/components/icloud/device_tracker.py
@@ -15,7 +15,7 @@ from .const import (
     DEVICE_LOCATION_LATITUDE,
     DEVICE_LOCATION_LONGITUDE,
     DOMAIN,
-    TRACKER_UPDATE,
+    SERVICE_UPDATE,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -77,11 +77,6 @@ class IcloudTrackerEntity(TrackerEntity):
         """Return longitude value of the device."""
         return self._device.location[DEVICE_LOCATION_LONGITUDE]
 
-    @property
-    def should_poll(self) -> bool:
-        """No polling needed."""
-        return False
-
     @property
     def battery_level(self) -> int:
         """Return the battery level of the device."""
@@ -112,10 +107,15 @@ class IcloudTrackerEntity(TrackerEntity):
             "model": self._device.device_model,
         }
 
+    @property
+    def should_poll(self) -> bool:
+        """No polling needed."""
+        return False
+
     async def async_added_to_hass(self):
         """Register state update callback."""
         self._unsub_dispatcher = async_dispatcher_connect(
-            self.hass, TRACKER_UPDATE, self.async_write_ha_state
+            self.hass, SERVICE_UPDATE, self.async_write_ha_state
         )
 
     async def async_will_remove_from_hass(self):
diff --git a/homeassistant/components/icloud/sensor.py b/homeassistant/components/icloud/sensor.py
index 4351d4ffa19..f6c87ed12d0 100644
--- a/homeassistant/components/icloud/sensor.py
+++ b/homeassistant/components/icloud/sensor.py
@@ -4,12 +4,13 @@ from typing import Dict
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_USERNAME, DEVICE_CLASS_BATTERY
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.icon import icon_for_battery_level
 from homeassistant.helpers.typing import HomeAssistantType
 
 from . import IcloudDevice
-from .const import DOMAIN
+from .const import DOMAIN, SERVICE_UPDATE
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -35,6 +36,7 @@ class IcloudDeviceBatterySensor(Entity):
     def __init__(self, device: IcloudDevice):
         """Initialize the battery sensor."""
         self._device = device
+        self._unsub_dispatcher = None
 
     @property
     def unique_id(self) -> str:
@@ -83,3 +85,18 @@ class IcloudDeviceBatterySensor(Entity):
             "manufacturer": "Apple",
             "model": self._device.device_model,
         }
+
+    @property
+    def should_poll(self) -> bool:
+        """No polling needed."""
+        return False
+
+    async def async_added_to_hass(self):
+        """Register state update callback."""
+        self._unsub_dispatcher = async_dispatcher_connect(
+            self.hass, SERVICE_UPDATE, self.async_write_ha_state
+        )
+
+    async def async_will_remove_from_hass(self):
+        """Clean up after entity before removal."""
+        self._unsub_dispatcher()