Use dispatcher for unifi heartbeat tracking (#45211)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
J. Nick Koston 2021-01-16 16:10:52 -10:00 committed by GitHub
parent b71a9b5e28
commit 41e7d960ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 62 additions and 38 deletions

View File

@ -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

View File

@ -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()

View File

@ -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)