From 41e7d960ee0a77a808f7dc3a4cebc1f07ed075ac Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 16 Jan 2021 16:10:52 -1000 Subject: [PATCH] Use dispatcher for unifi heartbeat tracking (#45211) Co-authored-by: Martin Hjelmare --- homeassistant/components/unifi/controller.py | 55 +++++++++++-------- .../components/unifi/device_tracker.py | 44 ++++++++++----- tests/components/unifi/test_device_tracker.py | 1 + 3 files changed, 62 insertions(+), 38 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 9f264e6cd1c..095eda9dff9 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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._watch_disconnected_entites = [] + self._cancel_heartbeat_check = None + 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.hass, self._async_check_for_disconnected, CHECK_DISCONNECTED_INTERVAL + self._cancel_heartbeat_check = async_track_time_interval( + self.hass, self._async_check_for_stale, CHECK_HEARTBEAT_INTERVAL ) return True @callback - def add_disconnected_check(self, entity: Entity) -> None: - """Add an entity to watch for disconnection.""" - self._watch_disconnected_entites.append(entity) + def async_heartbeat( + self, unique_id: str, heartbeat_expire_time: Optional[datetime] = None + ) -> 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: - """Remove an entity to watch for disconnection.""" - self._watch_disconnected_entites.remove(entity) - - @callback - def _async_check_for_disconnected(self, *_) -> None: + def _async_check_for_stale(self, *_) -> None: """Check for any devices scheduled to be marked disconnected.""" now = dt_util.utcnow() - for entity in self._watch_disconnected_entites: - disconnected_time = entity.disconnected_time - if disconnected_time is not None and now > disconnected_time: - entity.make_disconnected() + for unique_id, heartbeat_expire_time in self._heartbeat_time.items(): + if now > heartbeat_expire_time: + async_dispatcher_send( + 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: - self._cancel_disconnected_check() - self._cancel_disconnected_check = None + if self._cancel_heartbeat_check: + self._cancel_heartbeat_check() + self._cancel_heartbeat_check = None return True diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 22e4904ab8b..2eca5de5725 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/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 = ( - dt_util.utcnow() + self.controller.option_detection_time + self.controller.async_heartbeat( + self.unique_id, dt_util.utcnow() + self.controller.option_detection_time ) + self.heartbeat_check = True super().async_update_callback() @callback - def make_disconnected(self, *_): - """Mark client as disconnected.""" + def _make_disconnected(self, *_): + """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( - seconds=self.device.next_interval + 60 + self.controller.async_heartbeat( + 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() diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 6936b2e0fb5..8890ae7e959 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/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)