UniFi - Make clients proper push based (#35273)

* Improve client tracker to be more comprehensible and streamlined
Improve block switches
Improve tests

* Small clean up

* Add descriptions on ssid test

* Improve test

* Make polling default off, only POE clients left to verify

* Minor improvements

* On removal cancel scheduled updates

* POE works without polling now

* Combine else and if to an elif

homeassistant/components/unifi/device_tracker.py

@@ -2,7 +2,13 @@
 from datetime import timedelta
 import logging
 
-from aiounifi.api import SOURCE_DATA
+from aiounifi.api import SOURCE_DATA, SOURCE_EVENT
+from aiounifi.events import (
+    WIRED_CLIENT_CONNECTED,
+    WIRELESS_CLIENT_CONNECTED,
+    WIRELESS_CLIENT_ROAM,
+    WIRELESS_CLIENT_ROAMRADIO,
+)
 
 from homeassistant.components.device_tracker import DOMAIN
 from homeassistant.components.device_tracker.config_entry import ScannerEntity
@@ -19,6 +25,9 @@ from .unifi_entity_base import UniFiBase
 
 LOGGER = logging.getLogger(__name__)
 
+CLIENT_TRACKER = "client"
+DEVICE_TRACKER = "device"
+
 CLIENT_CONNECTED_ATTRIBUTES = [
     "_is_guest_by_uap",
     "ap_mac",
@@ -41,8 +50,12 @@ CLIENT_STATIC_ATTRIBUTES = [
     "oui",
 ]
 
-CLIENT_TRACKER = "client"
+WIRED_CONNECTION = (WIRED_CLIENT_CONNECTED,)
-DEVICE_TRACKER = "device"
+WIRELESS_CONNECTION = (
+    WIRELESS_CLIENT_CONNECTED,
+    WIRELESS_CLIENT_ROAM,
+    WIRELESS_CLIENT_ROAMRADIO,
+)
 
 
 async def async_setup_entry(hass, config_entry, async_add_entities):
@@ -120,79 +133,83 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         """Set up tracked client."""
         super().__init__(client, controller)
 
+        self.schedule_update = False
         self.cancel_scheduled_update = None
-        self.is_disconnected = None
+        self._is_connected = False
-        self.wired_bug = None
+        if self.client.last_seen:
-        if self.is_wired != self.client.is_wired:
+            self._is_connected = (
-            self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time
+                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
+            )
+            if self._is_connected:
+                self.schedule_update = True
 
-    @property
+    async def async_will_remove_from_hass(self) -> None:
-    def is_connected(self):
+        """Disconnect object when removed."""
-        """Return true if the client is connected to the network.
+        if self.cancel_scheduled_update:
+            self.cancel_scheduled_update()
+        await super().async_will_remove_from_hass()
 
-        If connected to unwanted ssid return False.
+    @callback
-        If is_wired and client.is_wired differ it means that the device is offline and UniFi bug shows device as wired.
+    def async_update_callback(self) -> None:
-        """
+        """Update the clients state."""
 
         @callback
-        def _scheduled_update(now):
+        def _make_disconnected(now):
-            """Scheduled callback for update."""
+            """Mark client as disconnected."""
-            self.is_disconnected = True
+            self._is_connected = False
             self.cancel_scheduled_update = None
             self.async_write_ha_state()
 
-        if (self.is_wired and self.wired_connection) or (
+        if self.client.last_updated == SOURCE_EVENT:
-            not self.is_wired and self.wireless_connection
+
-        ):
+            if (self.is_wired and self.client.event.event in WIRED_CONNECTION) or (
+                not self.is_wired and self.client.event.event in WIRELESS_CONNECTION
+            ):
+                self._is_connected = True
+                self.schedule_update = False
+                if self.cancel_scheduled_update:
+                    self.cancel_scheduled_update()
+                    self.cancel_scheduled_update = None
+
+            # Ignore extra scheduled update from wired bug
+            elif not self.cancel_scheduled_update:
+                self.schedule_update = True
+
+        elif not self.client.event and self.client.last_updated == SOURCE_DATA:
+
+            if self.is_wired == self.client.is_wired:
+                self._is_connected = True
+                self.schedule_update = True
+
+        if self.schedule_update:
+            self.schedule_update = False
+
             if self.cancel_scheduled_update:
                 self.cancel_scheduled_update()
-                self.cancel_scheduled_update = None
 
-            self.is_disconnected = False
+            self.cancel_scheduled_update = async_track_point_in_utc_time(
+                self.hass,
+                _make_disconnected,
+                dt_util.utcnow() + self.controller.option_detection_time,
+            )
 
-        if (self.is_wired and self.wired_connection is False) or (
+        super().async_update_callback()
-            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,
-                )
 
-        # SSID filter
+    @property
+    def is_connected(self):
+        """Return true if the client is connected to the network."""
         if (
             not self.is_wired
             and self.client.essid
             and self.controller.option_ssid_filter
             and self.client.essid not in self.controller.option_ssid_filter
-            and not self.cancel_scheduled_update
         ):
             return False
 
-        # A client that has never been seen cannot be connected.
+        return self._is_connected
-        if self.client.last_seen is None:
-            return False
-
-        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()
-            since_last_seen = dt_util.utcnow() - self.wired_bug
-
-        else:
-            self.wired_bug = None
-
-            since_last_seen = dt_util.utcnow() - dt_util.utc_from_timestamp(
-                float(self.client.last_seen)
-            )
-
-        if since_last_seen < self.controller.option_detection_time:
-            return True
-
-        return False
 
     @property
     def source_type(self):
@@ -213,7 +230,7 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
 
         for variable in CLIENT_STATIC_ATTRIBUTES + CLIENT_CONNECTED_ATTRIBUTES:
             if variable in self.client.raw:
-                if self.is_disconnected and variable in CLIENT_CONNECTED_ATTRIBUTES:
+                if not self.is_connected and variable in CLIENT_CONNECTED_ATTRIBUTES:
                     continue
                 attributes[variable] = self.client.raw[variable]
 
@@ -227,12 +244,12 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         elif self.is_wired:
             if not self.controller.option_track_wired_clients:
                 await self.async_remove()
-        else:
+
-            if (
+        elif (
-                self.controller.option_ssid_filter
+            self.controller.option_ssid_filter
-                and self.client.essid not in self.controller.option_ssid_filter
+            and self.client.essid not in self.controller.option_ssid_filter
-            ):
+        ):
-                await self.async_remove()
+            await self.async_remove()
 
 
 class UniFiDeviceTracker(UniFiBase, ScannerEntity):
@@ -261,8 +278,10 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect device object when removed."""
-        await super().async_will_remove_from_hass()
         self.device.remove_callback(self.async_update_callback)
+        if self.cancel_scheduled_update:
+            self.cancel_scheduled_update()
+        await super().async_will_remove_from_hass()
 
     @callback
     def async_update_callback(self):
@@ -287,8 +306,7 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
                 dt_util.utcnow() + timedelta(seconds=self.device.next_interval + 10),
             )
 
-        LOGGER.debug("Updating device %s (%s)", self.entity_id, self.device.mac)
+        super().async_update_callback()
-        self.async_write_ha_state()
 
     @property
     def is_connected(self):
@@ -353,8 +371,3 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
         """Config entry options are updated, remove entity if option is disabled."""
         if not self.controller.option_track_devices:
             await self.async_remove()
-
-    @property
-    def should_poll(self) -> bool:
-        """No polling needed."""
-        return False

homeassistant/components/unifi/switch.py

@@ -1,6 +1,14 @@
 """Support for devices connected to UniFi POE."""
 import logging
 
+from aiounifi.api import SOURCE_EVENT
+from aiounifi.events import (
+    WIRED_CLIENT_BLOCKED,
+    WIRED_CLIENT_UNBLOCKED,
+    WIRELESS_CLIENT_BLOCKED,
+    WIRELESS_CLIENT_UNBLOCKED,
+)
+
 from homeassistant.components.switch import DOMAIN, SwitchEntity
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -14,6 +22,9 @@ LOGGER = logging.getLogger(__name__)
 BLOCK_SWITCH = "block"
 POE_SWITCH = "poe"
 
+CLIENT_BLOCKED = (WIRED_CLIENT_BLOCKED, WIRELESS_CLIENT_BLOCKED)
+CLIENT_UNBLOCKED = (WIRED_CLIENT_UNBLOCKED, WIRELESS_CLIENT_UNBLOCKED)
+
 
 async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
     """Component doesn't support configuration through configuration.yaml."""
@@ -237,10 +248,26 @@ class UniFiBlockClientSwitch(UniFiClient, SwitchEntity):
     DOMAIN = DOMAIN
     TYPE = BLOCK_SWITCH
 
+    def __init__(self, client, controller):
+        """Set up block switch."""
+        super().__init__(client, controller)
+
+        self._is_blocked = self.client.blocked
+
+    @callback
+    def async_update_callback(self) -> None:
+        """Update the clients state."""
+        if self.client.last_updated == SOURCE_EVENT:
+
+            if self.client.event.event in CLIENT_BLOCKED + CLIENT_UNBLOCKED:
+                self._is_blocked = self.client.event.event in CLIENT_BLOCKED
+
+        super().async_update_callback()
+
     @property
     def is_on(self):
         """Return true if client is allowed to connect."""
-        return not self.is_blocked
+        return not self._is_blocked
 
     async def async_turn_on(self, **kwargs):
         """Turn on connectivity for client."""
@@ -253,7 +280,7 @@ class UniFiBlockClientSwitch(UniFiClient, SwitchEntity):
     @property
     def icon(self):
         """Return the icon to use in the frontend."""
-        if self.is_blocked:
+        if self._is_blocked:
             return "mdi:network-off"
         return "mdi:network"
 

homeassistant/components/unifi/unifi_client.py

@@ -2,37 +2,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_ROAM,
-    WIRELESS_CLIENT_ROAMRADIO,
-    WIRELESS_CLIENT_UNBLOCKED,
-)
-
-from homeassistant.core import callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 
 from .unifi_entity_base import UniFiBase
 
 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,
-    WIRELESS_CLIENT_ROAM,
-    WIRELESS_CLIENT_ROAMRADIO,
-)
-
 
 class UniFiClient(UniFiBase):
     """Base class for UniFi clients."""
@@ -43,9 +18,6 @@ class UniFiClient(UniFiBase):
         super().__init__(controller)
 
         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
 
     @property
     def mac(self):
@@ -59,33 +31,8 @@ class UniFiClient(UniFiBase):
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect client object when removed."""
-        await super().async_will_remove_from_hass()
         self.client.remove_callback(self.async_update_callback)
-
+        await super().async_will_remove_from_hass()
-    @callback
-    def async_update_callback(self) -> None:
-        """Update the clients state."""
-        if self._is_wired and self.client.mac in self.controller.wireless_clients:
-            self._is_wired = False
-
-        if self.client.last_updated == SOURCE_EVENT:
-            if self.client.event.event in WIRELESS_CLIENT:
-                self.wireless_connection = self.client.event.event in (
-                    WIRELESS_CLIENT_CONNECTED,
-                    WIRELESS_CLIENT_ROAM,
-                    WIRELESS_CLIENT_ROAMRADIO,
-                )
-
-            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)
-        self.async_write_ha_state()
 
     @property
     def is_wired(self):
@@ -93,6 +40,9 @@ class UniFiClient(UniFiBase):
 
         Allows disabling logic to keep track of clients affected by UniFi wired bug marking wireless devices as wired. This is useful when running a network not only containing UniFi APs.
         """
+        if self._is_wired and self.client.mac in self.controller.wireless_clients:
+            self._is_wired = False
+
         if self.controller.option_ignore_wired_bug:
             return self.client.is_wired
         return self._is_wired

homeassistant/components/unifi/unifi_entity_base.py

@@ -40,6 +40,7 @@ class UniFiBase(Entity):
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect object when removed."""
+        LOGGER.debug("Removing %s entity %s (%s)", self.TYPE, self.entity_id, self.mac)
         self.controller.entities[self.DOMAIN][self.TYPE].remove(self.mac)
 
     async def async_remove(self):
@@ -69,9 +70,10 @@ class UniFiBase(Entity):
         entity_registry.async_remove(self.entity_id)
 
     @callback
-    def async_update_callback(self):
+    def async_update_callback(self) -> None:
         """Update the entity's state."""
-        raise NotImplementedError
+        LOGGER.debug("Updating %s entity %s (%s)", self.TYPE, self.entity_id, self.mac)
+        self.async_write_ha_state()
 
     async def options_updated(self) -> None:
         """Config entry options are updated, remove entity if option is disabled."""
@@ -85,4 +87,4 @@ class UniFiBase(Entity):
     @property
     def should_poll(self) -> bool:
         """No polling needed."""
-        return True
+        return False

tests/components/unifi/test_device_tracker.py

@@ -29,7 +29,6 @@ import homeassistant.util.dt as dt_util
 
 from .test_controller import ENTRY_CONFIG, setup_unifi_integration
 
-from tests.async_mock import patch
 from tests.common import async_fire_time_changed
 
 CLIENT_1 = {
@@ -167,7 +166,6 @@ async def test_tracked_wireless_clients(hass):
 
     # State change signalling works without events
     client_1_copy = copy(CLIENT_1)
-    client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     controller.api.websocket._data = {
         "meta": {"message": MESSAGE_CLIENT},
         "data": [client_1_copy],
@@ -205,8 +203,6 @@ async def test_tracked_wireless_clients(hass):
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
 
-    # test wired bug
-
 
 async def test_tracked_clients(hass):
     """Test the update_items function with some clients."""
@@ -217,10 +213,9 @@ async def test_tracked_clients(hass):
         hass,
         options={CONF_SSID_FILTER: ["ssid"]},
         clients_response=[CLIENT_1, CLIENT_2, CLIENT_3, CLIENT_5, client_4_copy],
-        devices_response=[DEVICE_1, DEVICE_2],
+        known_wireless_clients=(CLIENT_4["mac"],),
-        known_wireless_clients=([CLIENT_4["mac"]]),
     )
-    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 6
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 4
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
@@ -246,9 +241,9 @@ async def test_tracked_clients(hass):
 
     # State change signalling works
     client_1_copy = copy(CLIENT_1)
-    client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
     controller.api.message_handler(event)
+    await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
@@ -357,7 +352,7 @@ async def test_controller_state_change(hass):
     await hass.async_block_till_done()
 
     client_1 = hass.states.get("device_tracker.client_1")
-    assert client_1.state == "not_home"
+    assert client_1.state == "home"
 
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1.state == "home"
@@ -499,29 +494,56 @@ async def test_option_track_devices(hass):
 
 
 async def test_option_ssid_filter(hass):
-    """Test the SSID filter works."""
+    """Test the SSID filter works.
-    controller = await setup_unifi_integration(hass, clients_response=[CLIENT_3])
+
-    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
+    Client 1 will travel from a supported SSID to an unsupported ssid.
+    Client 3 will be removed on change of options since it is in an unsupported SSID.
+    """
+    client_1_copy = copy(CLIENT_1)
+    client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+
+    controller = await setup_unifi_integration(
+        hass, clients_response=[client_1_copy, CLIENT_3]
+    )
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
 
     client_3 = hass.states.get("device_tracker.client_3")
     assert client_3
 
-    # Set SSID filter
+    # Setting SSID filter will remove clients outside of filter
     hass.config_entries.async_update_entry(
         controller.config_entry, options={CONF_SSID_FILTER: ["ssid"]},
     )
     await hass.async_block_till_done()
 
+    # Not affected by SSID filter
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
+    # Removed due to SSID filter
     client_3 = hass.states.get("device_tracker.client_3")
     assert not client_3
 
+    # Roams to SSID outside of filter
+    client_1_copy = copy(CLIENT_1)
+    client_1_copy["essid"] = "other_ssid"
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
+    controller.api.message_handler(event)
+    # Data update while SSID filter is in effect shouldn't create the client
     client_3_copy = copy(CLIENT_3)
     client_3_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
-    # SSID filter active even though time stamp should mark as home
+    # SSID filter marks client as away
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+
+    # SSID still outside of filter
     client_3 = hass.states.get("device_tracker.client_3")
     assert not client_3
 
@@ -529,13 +551,37 @@ async def test_option_ssid_filter(hass):
     hass.config_entries.async_update_entry(
         controller.config_entry, options={CONF_SSID_FILTER: []},
     )
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
+    controller.api.message_handler(event)
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
     client_3 = hass.states.get("device_tracker.client_3")
     assert client_3.state == "home"
 
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+
+    # Client won't go away until after next update
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3.state == "home"
+
+    # Trigger update to get client marked as away
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [CLIENT_3]}
+    controller.api.message_handler(event)
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
+    await hass.async_block_till_done()
+
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3.state == "not_home"
+
 
 async def test_wireless_client_go_wired_issue(hass):
     """Test the solution to catch wireless device go wired UniFi issue.
@@ -548,38 +594,49 @@ async def test_wireless_client_go_wired_issue(hass):
     controller = await setup_unifi_integration(hass, clients_response=[client_1_client])
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
 
+    # Client is wireless
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
 
+    # Trigger wired bug
     client_1_client["is_wired"] = True
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Wired bug fix keeps client marked as wireless
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
 
-    with patch.object(
+    # Pass time
-        dt_util, "utcnow", return_value=(dt_util.utcnow() + timedelta(minutes=5)),
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
-    ):
+    await hass.async_block_till_done()
-        event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
-        controller.api.message_handler(event)
-        await hass.async_block_till_done()
 
-        client_1 = hass.states.get("device_tracker.client_1")
+    # Marked as home according to the timer
-        assert client_1.state == "not_home"
+    client_1 = hass.states.get("device_tracker.client_1")
-        assert client_1.attributes["is_wired"] is False
+    assert client_1.state == "not_home"
+    assert client_1.attributes["is_wired"] is False
 
-    client_1_client["is_wired"] = False
+    # Try to mark client as connected
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
-    event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Make sure it don't go online again until wired bug disappears
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+    assert client_1.attributes["is_wired"] is False
+
+    # Make client wireless
+    client_1_client["is_wired"] = False
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    controller.api.message_handler(event)
+    await hass.async_block_till_done()
+
+    # Client is no longer affected by wired bug and can be marked online
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
@@ -595,27 +652,49 @@ async def test_option_ignore_wired_bug(hass):
     )
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
 
+    # Client is wireless
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
 
+    # Trigger wired bug
     client_1_client["is_wired"] = True
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
-    event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Wired bug in effect
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is True
 
-    client_1_client["is_wired"] = False
+    # pass time
-    client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    async_fire_time_changed(hass, dt_util.utcnow() + controller.option_detection_time)
-    event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
+    await hass.async_block_till_done()
+
+    # Timer marks client as away
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+    assert client_1.attributes["is_wired"] is True
+
+    # Mark client as connected again
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
+    # Ignoring wired bug allows client to go home again even while affected
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+    assert client_1.attributes["is_wired"] is True
+
+    # Make client wireless
+    client_1_client["is_wired"] = False
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
+    controller.api.message_handler(event)
+    await hass.async_block_till_done()
+
+    # Client is wireless and still connected
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1.state == "home"
     assert client_1.attributes["is_wired"] is False
@@ -665,7 +744,7 @@ async def test_restoring_client(hass):
 
 async def test_dont_track_clients(hass):
     """Test don't track clients config works."""
-    await setup_unifi_integration(
+    controller = await setup_unifi_integration(
         hass,
         options={CONF_TRACK_CLIENTS: False},
         clients_response=[CLIENT_1],
@@ -678,12 +757,24 @@ async def test_dont_track_clients(hass):
 
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1 is not None
-    assert device_1.state == "home"
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_CLIENTS: True},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
 
 
 async def test_dont_track_devices(hass):
     """Test don't track devices config works."""
-    await setup_unifi_integration(
+    controller = await setup_unifi_integration(
         hass,
         options={CONF_TRACK_DEVICES: False},
         clients_response=[CLIENT_1],
@@ -693,15 +784,27 @@ async def test_dont_track_devices(hass):
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
-    assert client_1.state == "not_home"
 
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1 is None
 
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_DEVICES: True},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
 
 async def test_dont_track_wired_clients(hass):
     """Test don't track wired clients config works."""
-    await setup_unifi_integration(
+    controller = await setup_unifi_integration(
         hass,
         options={CONF_TRACK_WIRED_CLIENTS: False},
         clients_response=[CLIENT_1, CLIENT_2],
@@ -710,7 +813,19 @@ async def test_dont_track_wired_clients(hass):
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
-    assert client_1.state == "not_home"
 
-    client_2 = hass.states.get("device_tracker.client_2")
+    client_2 = hass.states.get("device_tracker.wired_client")
     assert client_2 is None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_WIRED_CLIENTS: True},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None

tests/components/unifi/test_switch.py

@@ -211,6 +211,30 @@ EVENT_BLOCKED_CLIENT_CONNECTED = {
     "_id": "5ea331fa30c49e00f90ddc1a",
 }
 
+EVENT_BLOCKED_CLIENT_BLOCKED = {
+    "user": BLOCKED["mac"],
+    "hostname": BLOCKED["hostname"],
+    "key": "EVT_WC_Blocked",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587753456179,
+    "datetime": "2020-04-24T18:37:36Z",
+    "msg": f'User{[BLOCKED["mac"]]} has been blocked."',
+    "_id": "5ea331fa30c49e00f90ddc1a",
+}
+
+EVENT_BLOCKED_CLIENT_UNBLOCKED = {
+    "user": BLOCKED["mac"],
+    "hostname": BLOCKED["hostname"],
+    "key": "EVT_WC_Unblocked",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587753456179,
+    "datetime": "2020-04-24T18:37:36Z",
+    "msg": f'User{[BLOCKED["mac"]]} has been unblocked."',
+    "_id": "5ea331fa30c49e00f90ddc1a",
+}
+
 
 EVENT_CLIENT_2_CONNECTED = {
     "user": CLIENT_2["mac"],
@@ -368,6 +392,74 @@ async def test_remove_switches(hass):
     assert block_switch is None
 
 
+async def test_block_switches(hass):
+    """Test the update_items function with some clients."""
+    controller = await setup_unifi_integration(
+        hass,
+        options={
+            CONF_BLOCK_CLIENT: [BLOCKED["mac"], UNBLOCKED["mac"]],
+            CONF_TRACK_CLIENTS: False,
+            CONF_TRACK_DEVICES: False,
+        },
+        clients_response=[UNBLOCKED],
+        clients_all_response=[BLOCKED],
+    )
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+
+    blocked = hass.states.get("switch.block_client_1")
+    assert blocked is not None
+    assert blocked.state == "off"
+
+    unblocked = hass.states.get("switch.block_client_2")
+    assert unblocked is not None
+    assert unblocked.state == "on"
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_BLOCKED_CLIENT_UNBLOCKED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+    blocked = hass.states.get("switch.block_client_1")
+    assert blocked is not None
+    assert blocked.state == "on"
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_BLOCKED_CLIENT_BLOCKED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+    blocked = hass.states.get("switch.block_client_1")
+    assert blocked is not None
+    assert blocked.state == "off"
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
+    )
+    assert len(controller.mock_requests) == 5
+    assert controller.mock_requests[4] == {
+        "json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"},
+        "method": "post",
+        "path": "/cmd/stamgr",
+    }
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
+    )
+    assert len(controller.mock_requests) == 6
+    assert controller.mock_requests[5] == {
+        "json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
+        "method": "post",
+        "path": "/cmd/stamgr",
+    }
+
+
 async def test_new_client_discovered_on_block_control(hass):
     """Test if 2nd update has a new client."""
     controller = await setup_unifi_integration(