UniFi - Try to handle when UniFi erroneously marks offline client as wired (#26960)

* Add controls to catch when client goes offline and UniFi bug marks client as wired
* Device trackers shouldn't jump between going away and home
* POE control shouldn't add normally wireless clients as POE control switches

homeassistant/components/unifi/__init__.py

@@ -1,14 +1,13 @@
 """Support for devices connected to UniFi POE."""
 import voluptuous as vol
 
-from homeassistant.components.unifi.config_flow import (
-    get_controller_id_from_config_entry,
-)
 from homeassistant.const import CONF_HOST
+from homeassistant.core import callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 
 import homeassistant.helpers.config_validation as cv
 
+from .config_flow import get_controller_id_from_config_entry
 from .const import (
     ATTR_MANUFACTURER,
     CONF_BLOCK_CLIENT,
@@ -20,9 +19,14 @@ from .const import (
     CONF_SSID_FILTER,
     DOMAIN,
     UNIFI_CONFIG,
+    UNIFI_WIRELESS_CLIENTS,
 )
 from .controller import UniFiController
 
+SAVE_DELAY = 10
+STORAGE_KEY = "unifi_data"
+STORAGE_VERSION = 1
+
 CONF_CONTROLLERS = "controllers"
 
 CONTROLLER_SCHEMA = vol.Schema(
@@ -61,6 +65,9 @@ async def async_setup(hass, config):
     if DOMAIN in config:
         hass.data[UNIFI_CONFIG] = config[DOMAIN][CONF_CONTROLLERS]
 
+    hass.data[UNIFI_WIRELESS_CLIENTS] = wireless_clients = UnifiWirelessClients(hass)
+    await wireless_clients.async_load()
+
     return True
 
 
@@ -70,9 +77,7 @@ async def async_setup_entry(hass, config_entry):
         hass.data[DOMAIN] = {}
 
     controller = UniFiController(hass, config_entry)
-
     controller_id = get_controller_id_from_config_entry(config_entry)
-
     hass.data[DOMAIN][controller_id] = controller
 
     if not await controller.async_setup():
@@ -99,3 +104,43 @@ async def async_unload_entry(hass, config_entry):
     controller_id = get_controller_id_from_config_entry(config_entry)
     controller = hass.data[DOMAIN].pop(controller_id)
     return await controller.async_reset()
+
+
+class UnifiWirelessClients:
+    """Class to store clients known to be wireless.
+
+    This is needed since wireless devices going offline might get marked as wired by UniFi.
+    """
+
+    def __init__(self, hass):
+        """Set up client storage."""
+        self.hass = hass
+        self.data = {}
+        self._store = hass.helpers.storage.Store(STORAGE_VERSION, STORAGE_KEY)
+
+    async def async_load(self):
+        """Load data from file."""
+        data = await self._store.async_load()
+
+        if data is not None:
+            self.data = data
+
+    @callback
+    def get_data(self, config_entry):
+        """Get data related to a specific controller."""
+        controller_id = get_controller_id_from_config_entry(config_entry)
+        data = self.data.get(controller_id, {"wireless_devices": []})
+        return set(data["wireless_devices"])
+
+    @callback
+    def update_data(self, data, config_entry):
+        """Update data and schedule to save to file."""
+        controller_id = get_controller_id_from_config_entry(config_entry)
+        self.data[controller_id] = {"wireless_devices": list(data)}
+
+        self._store.async_delay_save(self._data_to_save, SAVE_DELAY)
+
+    @callback
+    def _data_to_save(self):
+        """Return data of UniFi wireless clients to store in a file."""
+        return self.data

homeassistant/components/unifi/const.py

@@ -10,6 +10,7 @@ CONF_CONTROLLER = "controller"
 CONF_SITE_ID = "site"
 
 UNIFI_CONFIG = "unifi_config"
+UNIFI_WIRELESS_CLIENTS = "unifi_wireless_clients"
 
 CONF_BLOCK_CLIENT = "block_client"
 CONF_DETECTION_TIME = "detection_time"

homeassistant/components/unifi/controller.py

@@ -36,6 +36,7 @@ from .const import (
     DOMAIN,
     LOGGER,
     UNIFI_CONFIG,
+    UNIFI_WIRELESS_CLIENTS,
 )
 from .errors import AuthenticationRequired, CannotConnect
 
@@ -50,6 +51,7 @@ class UniFiController:
         self.available = True
         self.api = None
         self.progress = None
+        self.wireless_clients = None
 
         self._site_name = None
         self._site_role = None
@@ -128,6 +130,22 @@ class UniFiController:
         """Event specific per UniFi entry to signal new options."""
         return f"unifi-options-{CONTROLLER_ID.format(host=self.host, site=self.site)}"
 
+    def update_wireless_clients(self):
+        """Update set of known to be wireless clients."""
+        new_wireless_clients = set()
+
+        for client_id in self.api.clients:
+            if (
+                client_id not in self.wireless_clients
+                and not self.api.clients[client_id].is_wired
+            ):
+                new_wireless_clients.add(client_id)
+
+        if new_wireless_clients:
+            self.wireless_clients |= new_wireless_clients
+            unifi_wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
+            unifi_wireless_clients.update_data(self.wireless_clients, self.config_entry)
+
     async def request_update(self):
         """Request an update."""
         if self.progress is not None:
@@ -170,6 +188,8 @@ class UniFiController:
             LOGGER.info("Reconnected to controller %s", self.host)
             self.available = True
 
+        self.update_wireless_clients()
+
         async_dispatcher_send(self.hass, self.signal_update)
 
     async def async_setup(self):
@@ -197,6 +217,10 @@ class UniFiController:
             LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
             return False
 
+        wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
+        self.wireless_clients = wireless_clients.get_data(self.config_entry)
+        self.update_wireless_clients()
+
         self.import_configuration()
 
         self.config_entry.add_update_listener(self.async_options_updated)

homeassistant/components/unifi/device_tracker.py

@@ -26,7 +26,6 @@ DEVICE_ATTRIBUTES = [
     "ip",
     "is_11r",
     "is_guest",
-    "is_wired",
     "mac",
     "name",
     "noted",
@@ -121,6 +120,7 @@ class UniFiClientTracker(ScannerEntity):
         """Set up tracked client."""
         self.client = client
         self.controller = controller
+        self.is_wired = self.client.mac not in controller.wireless_clients
 
     @property
     def entity_registry_enabled_default(self):
@@ -129,13 +129,13 @@ class UniFiClientTracker(ScannerEntity):
             return False
 
         if (
-            not self.client.is_wired
+            not self.is_wired
             and self.controller.option_ssid_filter
             and self.client.essid not in self.controller.option_ssid_filter
         ):
             return False
 
-        if not self.controller.option_track_wired_clients and self.client.is_wired:
+        if not self.controller.option_track_wired_clients and self.is_wired:
             return False
 
         return True
@@ -145,18 +145,31 @@ class UniFiClientTracker(ScannerEntity):
         LOGGER.debug("New UniFi client tracker %s (%s)", self.name, self.client.mac)
 
     async def async_update(self):
-        """Synchronize state with controller."""
+        """Synchronize state with controller.
+
+        Make sure to update self.is_wired if client is wireless, there is an issue when clients go offline that they get marked as wired.
+        """
         LOGGER.debug(
             "Updating UniFi tracked client %s (%s)", self.entity_id, self.client.mac
         )
         await self.controller.request_update()
 
+        if self.is_wired and self.client.mac in self.controller.wireless_clients:
+            self.is_wired = False
+
     @property
     def is_connected(self):
-        """Return true if the client is connected to the network."""
+        """Return true if the client is connected to the network.
-        if (
+
-            dt_util.utcnow() - dt_util.utc_from_timestamp(float(self.client.last_seen))
+        If is_wired and client.is_wired differ it means that the device is offline and UniFi bug shows device as wired.
-        ) < self.controller.option_detection_time:
+        """
+        if self.is_wired == self.client.is_wired and (
+            (
+                dt_util.utcnow()
+                - dt_util.utc_from_timestamp(float(self.client.last_seen))
+            )
+            < self.controller.option_detection_time
+        ):
             return True
 
         return False
@@ -195,6 +208,8 @@ class UniFiClientTracker(ScannerEntity):
             if variable in self.client.raw:
                 attributes[variable] = self.client.raw[variable]
 
+        attributes["is_wired"] = self.is_wired
+
         return attributes
 
 

homeassistant/components/unifi/switch.py

@@ -88,7 +88,7 @@ def update_items(controller, async_add_entities, switches, switches_off):
         new_switches.append(switches[block_client_id])
         LOGGER.debug("New UniFi Block switch %s (%s)", client.hostname, client.mac)
 
-    # control poe
+    # control POE
     for client_id in controller.api.clients:
 
         poe_client_id = f"poe-{client_id}"
@@ -108,9 +108,10 @@ def update_items(controller, async_add_entities, switches, switches_off):
             pass
         # Network device with active POE
         elif (
-            not client.is_wired
+            client_id in controller.wireless_clients
             or client.sw_mac not in devices
             or not devices[client.sw_mac].ports[client.sw_port].port_poe
+            or not devices[client.sw_mac].ports[client.sw_port].poe_enable
             or controller.mac == client.mac
         ):
             continue

tests/components/unifi/test_controller.py

@@ -8,6 +8,7 @@ from homeassistant.components.unifi.const import (
     CONF_CONTROLLER,
     CONF_SITE_ID,
     UNIFI_CONFIG,
+    UNIFI_WIRELESS_CLIENTS,
 )
 from homeassistant.const import (
     CONF_HOST,
@@ -49,7 +50,8 @@ async def test_controller_setup():
                 controller.CONF_DETECTION_TIME: 30,
                 controller.CONF_SSID_FILTER: ["ssid"],
             }
-        ]
+        ],
+        UNIFI_WIRELESS_CLIENTS: Mock(),
     }
     entry = Mock()
     entry.data = ENTRY_CONFIG
@@ -57,6 +59,7 @@ async def test_controller_setup():
     api = Mock()
     api.initialize.return_value = mock_coro(True)
     api.sites.return_value = mock_coro(CONTROLLER_SITES)
+    api.clients = []
 
     unifi_controller = controller.UniFiController(hass, entry)
 
@@ -100,7 +103,8 @@ async def test_controller_site():
 async def test_controller_mac():
     """Test that it is possible to identify controller mac."""
     hass = Mock()
-    hass.data = {UNIFI_CONFIG: {}}
+    hass.data = {UNIFI_CONFIG: {}, UNIFI_WIRELESS_CLIENTS: Mock()}
+    hass.data[UNIFI_WIRELESS_CLIENTS].get_data.return_value = set()
     entry = Mock()
     entry.data = ENTRY_CONFIG
     entry.options = {}
@@ -123,7 +127,7 @@ async def test_controller_mac():
 async def test_controller_no_mac():
     """Test that it works to not find the controllers mac."""
     hass = Mock()
-    hass.data = {UNIFI_CONFIG: {}}
+    hass.data = {UNIFI_CONFIG: {}, UNIFI_WIRELESS_CLIENTS: Mock()}
     entry = Mock()
     entry.data = ENTRY_CONFIG
     entry.options = {}
@@ -133,6 +137,7 @@ async def test_controller_no_mac():
     api.initialize.return_value = mock_coro(True)
     api.clients = {"client1": client}
     api.sites.return_value = mock_coro(CONTROLLER_SITES)
+    api.clients = {}
 
     unifi_controller = controller.UniFiController(hass, entry)
 
@@ -195,13 +200,14 @@ async def test_reset_if_entry_had_wrong_auth():
 async def test_reset_unloads_entry_if_setup():
     """Calling reset when the entry has been setup."""
     hass = Mock()
-    hass.data = {UNIFI_CONFIG: {}}
+    hass.data = {UNIFI_CONFIG: {}, UNIFI_WIRELESS_CLIENTS: Mock()}
     entry = Mock()
     entry.data = ENTRY_CONFIG
     entry.options = {}
     api = Mock()
     api.initialize.return_value = mock_coro(True)
     api.sites.return_value = mock_coro(CONTROLLER_SITES)
+    api.clients = []
 
     unifi_controller = controller.UniFiController(hass, entry)
 

tests/components/unifi/test_device_tracker.py

@@ -1,9 +1,11 @@
 """The tests for the UniFi device tracker platform."""
 from collections import deque
 from copy import copy
-from unittest.mock import Mock
+
 from datetime import timedelta
 
+from asynctest import Mock
+
 import pytest
 
 from aiounifi.clients import Clients, ClientsAll
@@ -19,6 +21,7 @@ from homeassistant.components.unifi.const import (
     CONF_TRACK_WIRED_CLIENTS,
     CONTROLLER_ID as CONF_CONTROLLER_ID,
     UNIFI_CONFIG,
+    UNIFI_WIRELESS_CLIENTS,
 )
 from homeassistant.const import (
     CONF_HOST,
@@ -96,7 +99,7 @@ CONTROLLER_DATA = {
     CONF_PASSWORD: "mock-pswd",
     CONF_PORT: 1234,
     CONF_SITE_ID: "mock-site",
-    CONF_VERIFY_SSL: True,
+    CONF_VERIFY_SSL: False,
 }
 
 ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
@@ -108,7 +111,9 @@ CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site")
 def mock_controller(hass):
     """Mock a UniFi Controller."""
     hass.data[UNIFI_CONFIG] = {}
+    hass.data[UNIFI_WIRELESS_CLIENTS] = Mock()
     controller = unifi.UniFiController(hass, None)
+    controller.wireless_clients = set()
 
     controller.api = Mock()
     controller.mock_requests = []
@@ -253,6 +258,45 @@ async def test_tracked_devices(hass, mock_controller):
     assert device_1 is None
 
 
+async def test_wireless_client_go_wired_issue(hass, mock_controller):
+    """Test the solution to catch wireless device go wired UniFi issue.
+
+    UniFi has a known issue that when a wireless device goes away it sometimes gets marked as wired.
+    """
+    client_1_client = copy(CLIENT_1)
+    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    mock_controller.mock_client_responses.append([client_1_client])
+    mock_controller.mock_device_responses.append({})
+
+    await setup_controller(hass, mock_controller)
+    assert len(mock_controller.mock_requests) == 2
+    assert len(hass.states.async_all()) == 3
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+    assert client_1.state == "home"
+
+    client_1_client["is_wired"] = True
+    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    mock_controller.mock_client_responses.append([client_1_client])
+    mock_controller.mock_device_responses.append({})
+    await mock_controller.async_update()
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+
+    client_1_client["is_wired"] = False
+    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    mock_controller.mock_client_responses.append([client_1_client])
+    mock_controller.mock_device_responses.append({})
+    await mock_controller.async_update()
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
+
 async def test_restoring_client(hass, mock_controller):
     """Test the update_items function with some clients."""
     mock_controller.mock_client_responses.append([CLIENT_2])

tests/components/unifi/test_switch.py

@@ -17,6 +17,7 @@ from homeassistant.components.unifi.const import (
     CONF_SITE_ID,
     CONTROLLER_ID as CONF_CONTROLLER_ID,
     UNIFI_CONFIG,
+    UNIFI_WIRELESS_CLIENTS,
 )
 from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
@@ -221,7 +222,9 @@ CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site")
 def mock_controller(hass):
     """Mock a UniFi Controller."""
     hass.data[UNIFI_CONFIG] = {}
+    hass.data[UNIFI_WIRELESS_CLIENTS] = Mock()
     controller = unifi.UniFiController(hass, None)
+    controller.wireless_clients = set()
 
     controller._site_role = "admin"
 
@@ -326,7 +329,7 @@ async def test_switches(hass, mock_controller):
 
     await setup_controller(hass, mock_controller, options)
     assert len(mock_controller.mock_requests) == 3
-    assert len(hass.states.async_all()) == 5
+    assert len(hass.states.async_all()) == 4
 
     switch_1 = hass.states.get("switch.poe_client_1")
     assert switch_1 is not None