From 015e779d56f6ad13d5528dff948ed7f498348dc5 Mon Sep 17 00:00:00 2001
From: Robert Svensson <Kane610@users.noreply.github.com>
Date: Wed, 11 Mar 2020 17:24:52 +0100
Subject: [PATCH] UniFi - Client tracker schedules update on disconnect event
 (#32655)

---
 .../components/unifi/device_tracker.py        | 33 +++++++++++++++++++
 .../components/unifi/unifi_client.py          | 20 ++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py
index b398dad488b..e5d3bcfa82b 100644
--- a/homeassistant/components/unifi/device_tracker.py
+++ b/homeassistant/components/unifi/device_tracker.py
@@ -8,6 +8,7 @@ from homeassistant.components.unifi.config_flow import get_controller_from_confi
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.event import async_track_point_in_utc_time
 import homeassistant.util.dt as dt_util
 
 from .const import ATTR_MANUFACTURER
@@ -175,6 +176,8 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         """Set up tracked client."""
         super().__init__(client, controller)
 
+        self.cancel_scheduled_update = None
+        self.is_disconnected = None
         self.wired_bug = None
         if self.is_wired != self.client.is_wired:
             self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time
@@ -186,6 +189,14 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         If connected to unwanted ssid return False.
         If is_wired and client.is_wired differ it means that the device is offline and UniFi bug shows device as wired.
         """
+
+        @callback
+        def _scheduled_update(now):
+            """Scheduled callback for update."""
+            self.is_disconnected = True
+            self.cancel_scheduled_update = None
+            self.async_schedule_update_ha_state()
+
         if (
             not self.is_wired
             and self.controller.option_ssid_filter
@@ -193,6 +204,28 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         ):
             return False
 
+        if (self.is_wired and self.wired_connection) or (
+            not self.is_wired and self.wireless_connection
+        ):
+            if self.cancel_scheduled_update:
+                self.cancel_scheduled_update()
+                self.cancel_scheduled_update = None
+
+            self.is_disconnected = False
+
+        if (self.is_wired and self.wired_connection is False) or (
+            not self.is_wired and self.wireless_connection is False
+        ):
+            if not self.is_disconnected and not self.cancel_scheduled_update:
+                self.cancel_scheduled_update = async_track_point_in_utc_time(
+                    self.hass,
+                    _scheduled_update,
+                    dt_util.utcnow() + self.controller.option_detection_time,
+                )
+
+        if self.is_disconnected is not None:
+            return not self.is_disconnected
+
         if self.is_wired != self.client.is_wired:
             if not self.wired_bug:
                 self.wired_bug = dt_util.utcnow()
diff --git a/homeassistant/components/unifi/unifi_client.py b/homeassistant/components/unifi/unifi_client.py
index 4c1ce402e7f..b46771e574b 100644
--- a/homeassistant/components/unifi/unifi_client.py
+++ b/homeassistant/components/unifi/unifi_client.py
@@ -5,8 +5,12 @@ import logging
 from aiounifi.api import SOURCE_EVENT
 from aiounifi.events import (
     WIRED_CLIENT_BLOCKED,
+    WIRED_CLIENT_CONNECTED,
+    WIRED_CLIENT_DISCONNECTED,
     WIRED_CLIENT_UNBLOCKED,
     WIRELESS_CLIENT_BLOCKED,
+    WIRELESS_CLIENT_CONNECTED,
+    WIRELESS_CLIENT_DISCONNECTED,
     WIRELESS_CLIENT_UNBLOCKED,
 )
 
@@ -19,6 +23,8 @@ LOGGER = logging.getLogger(__name__)
 
 CLIENT_BLOCKED = (WIRED_CLIENT_BLOCKED, WIRELESS_CLIENT_BLOCKED)
 CLIENT_UNBLOCKED = (WIRED_CLIENT_UNBLOCKED, WIRELESS_CLIENT_UNBLOCKED)
+WIRED_CLIENT = (WIRED_CLIENT_CONNECTED, WIRED_CLIENT_DISCONNECTED)
+WIRELESS_CLIENT = (WIRELESS_CLIENT_CONNECTED, WIRELESS_CLIENT_DISCONNECTED)
 
 
 class UniFiClient(Entity):
@@ -32,6 +38,8 @@ class UniFiClient(Entity):
 
         self.is_wired = self.client.mac not in controller.wireless_clients
         self.is_blocked = self.client.blocked
+        self.wired_connection = None
+        self.wireless_connection = None
 
     async def async_added_to_hass(self) -> None:
         """Client entity created."""
@@ -57,7 +65,17 @@ class UniFiClient(Entity):
 
         if self.client.last_updated == SOURCE_EVENT:
 
-            if self.client.event.event in CLIENT_BLOCKED + CLIENT_UNBLOCKED:
+            if self.client.event.event in WIRELESS_CLIENT:
+                self.wireless_connection = (
+                    self.client.event.event == WIRELESS_CLIENT_CONNECTED
+                )
+
+            elif self.client.event.event in WIRED_CLIENT:
+                self.wired_connection = (
+                    self.client.event.event == WIRED_CLIENT_CONNECTED
+                )
+
+            elif self.client.event.event in CLIENT_BLOCKED + CLIENT_UNBLOCKED:
                 self.is_blocked = self.client.event.event in CLIENT_BLOCKED
 
         LOGGER.debug("Updating client %s %s", self.entity_id, self.client.mac)