Use dispatcher for unifi heartbeat tracking (#45211)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/unifi/controller.py

@@ -1,7 +1,8 @@
 """UniFi Controller abstraction."""
 import asyncio
-from datetime import timedelta
+from datetime import datetime, timedelta
 import ssl
+from typing import Optional
 
 from aiohttp import CookieJar
 import aiounifi
@@ -32,7 +33,6 @@ from homeassistant.core import callback
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client
 from homeassistant.helpers.dispatcher import async_dispatcher_send
-from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.event import async_track_time_interval
 import homeassistant.util.dt as dt_util
 
@@ -67,7 +67,7 @@ from .const import (
 from .errors import AuthenticationRequired, CannotConnect
 
 RETRY_TIMER = 15
-CHECK_DISCONNECTED_INTERVAL = timedelta(seconds=1)
+CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1)
 SUPPORTED_PLATFORMS = [TRACKER_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN]
 
 CLIENT_CONNECTED = (
@@ -98,8 +98,9 @@ class UniFiController:
         self._site_name = None
         self._site_role = None
 
-        self._cancel_disconnected_check = None
+        self._cancel_heartbeat_check = None
-        self._watch_disconnected_entites = []
+        self._heartbeat_dispatch = {}
+        self._heartbeat_time = {}
 
         self.entities = {}
 
@@ -298,6 +299,11 @@ class UniFiController:
         """Event specific per UniFi entry to signal new options."""
         return f"unifi-options-{self.controller_id}"
 
+    @property
+    def signal_heartbeat_missed(self):
+        """Event specific per UniFi device tracker to signal new heartbeat missed."""
+        return "unifi-heartbeat-missed"
+
     def update_wireless_clients(self):
         """Update set of known to be wireless clients."""
         new_wireless_clients = set()
@@ -382,31 +388,34 @@ class UniFiController:
 
         self.config_entry.add_update_listener(self.async_config_entry_updated)
 
-        self._cancel_disconnected_check = async_track_time_interval(
+        self._cancel_heartbeat_check = async_track_time_interval(
-            self.hass, self._async_check_for_disconnected, CHECK_DISCONNECTED_INTERVAL
+            self.hass, self._async_check_for_stale, CHECK_HEARTBEAT_INTERVAL
         )
 
         return True
 
     @callback
-    def add_disconnected_check(self, entity: Entity) -> None:
+    def async_heartbeat(
-        """Add an entity to watch for disconnection."""
+        self, unique_id: str, heartbeat_expire_time: Optional[datetime] = None
-        self._watch_disconnected_entites.append(entity)
+    ) -> None:
+        """Signal when a device has fresh home state."""
+        if heartbeat_expire_time is not None:
+            self._heartbeat_time[unique_id] = heartbeat_expire_time
+            return
+
+        if unique_id in self._heartbeat_time:
+            del self._heartbeat_time[unique_id]
 
     @callback
-    def remove_disconnected_check(self, entity: Entity) -> None:
+    def _async_check_for_stale(self, *_) -> None:
-        """Remove an entity to watch for disconnection."""
-        self._watch_disconnected_entites.remove(entity)
-
-    @callback
-    def _async_check_for_disconnected(self, *_) -> None:
         """Check for any devices scheduled to be marked disconnected."""
         now = dt_util.utcnow()
 
-        for entity in self._watch_disconnected_entites:
+        for unique_id, heartbeat_expire_time in self._heartbeat_time.items():
-            disconnected_time = entity.disconnected_time
+            if now > heartbeat_expire_time:
-            if disconnected_time is not None and now > disconnected_time:
+                async_dispatcher_send(
-                entity.make_disconnected()
+                    self.hass, f"{self.signal_heartbeat_missed}_{unique_id}"
+                )
 
     @staticmethod
     async def async_config_entry_updated(hass, config_entry) -> None:
@@ -461,9 +470,9 @@ class UniFiController:
             unsub_dispatcher()
         self.listeners = []
 
-        if self._cancel_disconnected_check:
+        if self._cancel_heartbeat_check:
-            self._cancel_disconnected_check()
+            self._cancel_heartbeat_check()
-            self._cancel_disconnected_check = None
+            self._cancel_heartbeat_check = None
 
         return True
 

homeassistant/components/unifi/device_tracker.py

@@ -143,8 +143,8 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         """Set up tracked client."""
         super().__init__(client, controller)
 
+        self.heartbeat_check = False
         self.schedule_update = False
-        self.disconnected_time = None
         self._is_connected = False
         if client.last_seen:
             self._is_connected = (
@@ -158,12 +158,18 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
 
     async def async_added_to_hass(self) -> None:
         """Watch object when added."""
-        self.controller.add_disconnected_check(self)
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                f"{self.controller.signal_heartbeat_missed}_{self.unique_id}",
+                self._make_disconnected,
+            )
+        )
         await super().async_added_to_hass()
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect object when removed."""
-        self.controller.remove_disconnected_check(self)
+        self.controller.async_heartbeat(self.unique_id)
         await super().async_will_remove_from_hass()
 
     @callback
@@ -176,10 +182,11 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
             ):
                 self._is_connected = True
                 self.schedule_update = False
-                self.disconnected_time = None
+                self.controller.async_heartbeat(self.unique_id)
+                self.heartbeat_check = False
 
             # Ignore extra scheduled update from wired bug
-            elif not self.disconnected_time:
+            elif not self.heartbeat_check:
                 self.schedule_update = True
 
         elif not self.client.event and self.client.last_updated == SOURCE_DATA:
@@ -189,15 +196,16 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
 
         if self.schedule_update:
             self.schedule_update = False
-            self.disconnected_time = (
+            self.controller.async_heartbeat(
-                dt_util.utcnow() + self.controller.option_detection_time
+                self.unique_id, dt_util.utcnow() + self.controller.option_detection_time
             )
+            self.heartbeat_check = True
 
         super().async_update_callback()
 
     @callback
-    def make_disconnected(self, *_):
+    def _make_disconnected(self, *_):
-        """Mark client as disconnected."""
+        """No heart beat by device."""
         self._is_connected = False
         self.async_write_ha_state()
 
@@ -282,7 +290,6 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
         super().__init__(device, controller)
 
         self._is_connected = device.state == 1
-        self.disconnected_time = None
 
     @property
     def device(self):
@@ -291,12 +298,18 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
 
     async def async_added_to_hass(self) -> None:
         """Watch object when added."""
-        self.controller.add_disconnected_check(self)
+        self.async_on_remove(
+            async_dispatcher_connect(
+                self.hass,
+                f"{self.controller.signal_heartbeat_missed}_{self.unique_id}",
+                self._make_disconnected,
+            )
+        )
         await super().async_added_to_hass()
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect object when removed."""
-        self.controller.remove_disconnected_check(self)
+        self.controller.async_heartbeat(self.unique_id)
         await super().async_will_remove_from_hass()
 
     @callback
@@ -305,8 +318,9 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
 
         if self.device.last_updated == SOURCE_DATA:
             self._is_connected = True
-            self.disconnected_time = dt_util.utcnow() + timedelta(
+            self.controller.async_heartbeat(
-                seconds=self.device.next_interval + 60
+                self.unique_id,
+                dt_util.utcnow() + timedelta(seconds=self.device.next_interval + 60),
             )
 
         elif (
@@ -319,7 +333,7 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
         super().async_update_callback()
 
     @callback
-    def make_disconnected(self, *_):
+    def _make_disconnected(self, *_):
         """No heart beat by device."""
         self._is_connected = False
         self.async_write_ha_state()

tests/components/unifi/test_device_tracker.py

@@ -633,6 +633,7 @@ async def test_option_ssid_filter(hass):
     # Trigger update to get client marked as away
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [CLIENT_3]}
     controller.api.message_handler(event)
+    await hass.async_block_till_done()
 
     new_time = (
         dt_util.utcnow() + controller.option_detection_time + timedelta(seconds=1)