UniFi - Improve signalling and handling of adding new entities (#34925)

* Change signalling and handling of adding devices

* Bump aiounifi to v19
Use the proper variable for roamradio

homeassistant/components/unifi/controller.py

@@ -6,14 +6,19 @@ import ssl
 from aiohttp import CookieJar
 import aiounifi
 from aiounifi.controller import (
-    DATA_CLIENT,
     DATA_CLIENT_REMOVED,
-    DATA_DEVICE,
     DATA_EVENT,
     SIGNAL_CONNECTION_STATE,
     SIGNAL_DATA,
 )
-from aiounifi.events import WIRELESS_CLIENT_CONNECTED, WIRELESS_GUEST_CONNECTED
+from aiounifi.events import (
+    ACCESS_POINT_CONNECTED,
+    GATEWAY_CONNECTED,
+    SWITCH_CONNECTED,
+    WIRED_CLIENT_CONNECTED,
+    WIRELESS_CLIENT_CONNECTED,
+    WIRELESS_GUEST_CONNECTED,
+)
 from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
 import async_timeout
 
@@ -55,6 +60,17 @@ from .errors import AuthenticationRequired, CannotConnect
 RETRY_TIMER = 15
 SUPPORTED_PLATFORMS = [TRACKER_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN]
 
+CLIENT_CONNECTED = (
+    WIRED_CLIENT_CONNECTED,
+    WIRELESS_CLIENT_CONNECTED,
+    WIRELESS_GUEST_CONNECTED,
+)
+DEVICE_CONNECTED = (
+    ACCESS_POINT_CONNECTED,
+    GATEWAY_CONNECTED,
+    SWITCH_CONNECTED,
+)
+
 
 class UniFiController:
     """Manages a single UniFi Controller."""
@@ -190,14 +206,33 @@ class UniFiController:
         elif signal == SIGNAL_DATA and data:
 
             if DATA_EVENT in data:
-                if next(iter(data[DATA_EVENT])).event in (
+                clients_connected = set()
-                    WIRELESS_CLIENT_CONNECTED,
+                devices_connected = set()
-                    WIRELESS_GUEST_CONNECTED,
+                wireless_clients_connected = False
-                ):
-                    self.update_wireless_clients()
 
-            elif DATA_CLIENT in data or DATA_DEVICE in data:
+                for event in data[DATA_EVENT]:
-                async_dispatcher_send(self.hass, self.signal_update)
+
+                    if event.event in CLIENT_CONNECTED:
+                        clients_connected.add(event.mac)
+
+                        if not wireless_clients_connected and event.event in (
+                            WIRELESS_CLIENT_CONNECTED,
+                            WIRELESS_GUEST_CONNECTED,
+                        ):
+                            wireless_clients_connected = True
+
+                    elif event.event in DEVICE_CONNECTED:
+                        devices_connected.add(event.mac)
+
+                if wireless_clients_connected:
+                    self.update_wireless_clients()
+                if clients_connected or devices_connected:
+                    async_dispatcher_send(
+                        self.hass,
+                        self.signal_update,
+                        clients_connected,
+                        devices_connected,
+                    )
 
             elif DATA_CLIENT_REMOVED in data:
                 async_dispatcher_send(

homeassistant/components/unifi/device_tracker.py

@@ -48,10 +48,15 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     controller.entities[DOMAIN] = {CLIENT_TRACKER: set(), DEVICE_TRACKER: set()}
 
     @callback
-    def items_added():
+    def items_added(
+        clients: set = controller.api.clients, devices: set = controller.api.devices
+    ) -> None:
         """Update the values of the controller."""
-        if controller.option_track_clients or controller.option_track_devices:
+        if controller.option_track_clients:
-            add_entities(controller, async_add_entities)
+            add_client_entities(controller, async_add_entities, clients)
+
+        if controller.option_track_devices:
+            add_device_entities(controller, async_add_entities, devices)
 
     for signal in (controller.signal_update, controller.signal_options_update):
         controller.listeners.append(async_dispatcher_connect(hass, signal, items_added))
@@ -60,38 +65,43 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
 
 @callback
-def add_entities(controller, async_add_entities):
+def add_client_entities(controller, async_add_entities, clients):
-    """Add new tracker entities from the controller."""
+    """Add new client tracker entities from the controller."""
     trackers = []
 
-    for items, tracker_class, track in (
+    for mac in clients:
-        (controller.api.clients, UniFiClientTracker, controller.option_track_clients),
+        if mac in controller.entities[DOMAIN][UniFiClientTracker.TYPE]:
-        (controller.api.devices, UniFiDeviceTracker, controller.option_track_devices),
-    ):
-        if not track:
             continue
 
-        for mac in items:
+        client = controller.api.clients[mac]
 
-            if mac in controller.entities[DOMAIN][tracker_class.TYPE]:
+        if mac not in controller.wireless_clients:
+            if not controller.option_track_wired_clients:
                 continue
+        elif (
+            client.essid
+            and controller.option_ssid_filter
+            and client.essid not in controller.option_ssid_filter
+        ):
+            continue
 
-            item = items[mac]
+        trackers.append(UniFiClientTracker(client, controller))
 
-            if tracker_class is UniFiClientTracker:
+    if trackers:
+        async_add_entities(trackers)
 
-                if mac not in controller.wireless_clients:
-                    if not controller.option_track_wired_clients:
-                        continue
-                else:
-                    if (
-                        item.essid
-                        and controller.option_ssid_filter
-                        and item.essid not in controller.option_ssid_filter
-                    ):
-                        continue
 
-            trackers.append(tracker_class(item, controller))
+@callback
+def add_device_entities(controller, async_add_entities, devices):
+    """Add new device tracker entities from the controller."""
+    trackers = []
+
+    for mac in devices:
+        if mac in controller.entities[DOMAIN][UniFiDeviceTracker.TYPE]:
+            continue
+
+        device = controller.api.devices[mac]
+        trackers.append(UniFiDeviceTracker(device, controller))
 
     if trackers:
         async_add_entities(trackers)
@@ -147,6 +157,7 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
                     dt_util.utcnow() + self.controller.option_detection_time,
                 )
 
+        # SSID filter
         if (
             not self.is_wired
             and self.client.essid
@@ -156,6 +167,10 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         ):
             return False
 
+        # A client that has never been seen cannot be connected.
+        if self.client.last_seen is None:
+            return False
+
         if self.is_disconnected is not None:
             return not self.is_disconnected
 
@@ -167,10 +182,6 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         else:
             self.wired_bug = None
 
-            # A client that has never been seen cannot be connected.
-            if self.client.last_seen is None:
-                return False
-
             since_last_seen = dt_util.utcnow() - dt_util.utc_from_timestamp(
                 float(self.client.last_seen)
             )
@@ -240,7 +251,6 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity):
     async def async_added_to_hass(self):
         """Subscribe to device events."""
         await super().async_added_to_hass()
-        LOGGER.debug("New device %s (%s)", self.entity_id, self.device.mac)
         self.device.register_callback(self.async_update_callback)
 
     async def async_will_remove_from_hass(self) -> None:

homeassistant/components/unifi/sensor.py

@@ -25,10 +25,12 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     controller.entities[DOMAIN] = {RX_SENSOR: set(), TX_SENSOR: set()}
 
     @callback
-    def items_added():
+    def items_added(
+        clients: set = controller.api.clients, devices: set = controller.api.devices
+    ) -> None:
         """Update the values of the controller."""
         if controller.option_allow_bandwidth_sensors:
-            add_entities(controller, async_add_entities)
+            add_entities(controller, async_add_entities, clients)
 
     for signal in (controller.signal_update, controller.signal_options_update):
         controller.listeners.append(async_dispatcher_connect(hass, signal, items_added))
@@ -37,14 +39,17 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
 
 @callback
-def add_entities(controller, async_add_entities):
+def add_entities(controller, async_add_entities, clients):
     """Add new sensor entities from the controller."""
     sensors = []
 
-    for mac in controller.api.clients:
+    for mac in clients:
         for sensor_class in (UniFiRxBandwidthSensor, UniFiTxBandwidthSensor):
-            if mac not in controller.entities[DOMAIN][sensor_class.TYPE]:
+            if mac in controller.entities[DOMAIN][sensor_class.TYPE]:
-                sensors.append(sensor_class(controller.api.clients[mac], controller))
+                continue
+
+            client = controller.api.clients[mac]
+            sensors.append(sensor_class(client, controller))
 
     if sensors:
         async_add_entities(sensors)

homeassistant/components/unifi/switch.py

@@ -51,10 +51,17 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
             controller.api.clients.process_raw([client.raw])
 
     @callback
-    def items_added():
+    def items_added(
+        clients: set = controller.api.clients, devices: set = controller.api.devices
+    ) -> None:
         """Update the values of the controller."""
-        if controller.option_block_clients or controller.option_poe_clients:
+        if controller.option_block_clients:
-            add_entities(controller, async_add_entities, previously_known_poe_clients)
+            add_block_entities(controller, async_add_entities, clients)
+
+        if controller.option_poe_clients:
+            add_poe_entities(
+                controller, async_add_entities, clients, previously_known_poe_clients
+            )
 
     for signal in (controller.signal_update, controller.signal_options_update):
         controller.listeners.append(async_dispatcher_connect(hass, signal, items_added))
@@ -64,62 +71,66 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
 
 @callback
-def add_entities(controller, async_add_entities, previously_known_poe_clients):
+def add_block_entities(controller, async_add_entities, clients):
     """Add new switch entities from the controller."""
     switches = []
 
     for mac in controller.option_block_clients:
-
+        if mac in controller.entities[DOMAIN][BLOCK_SWITCH] or mac not in clients:
-        if (
-            mac in controller.entities[DOMAIN][BLOCK_SWITCH]
-            or mac not in controller.api.clients
-        ):
             continue
 
         client = controller.api.clients[mac]
         switches.append(UniFiBlockClientSwitch(client, controller))
 
-    if controller.option_poe_clients:
+    if switches:
-        devices = controller.api.devices
+        async_add_entities(switches)
 
-        for mac in controller.api.clients:
 
-            poe_client_id = f"{POE_SWITCH}-{mac}"
+@callback
+def add_poe_entities(
+    controller, async_add_entities, clients, previously_known_poe_clients
+):
+    """Add new switch entities from the controller."""
+    switches = []
 
-            if mac in controller.entities[DOMAIN][POE_SWITCH]:
+    devices = controller.api.devices
-                continue
 
-            client = controller.api.clients[mac]
+    for mac in clients:
+        if mac in controller.entities[DOMAIN][POE_SWITCH]:
+            continue
 
-            if poe_client_id not in previously_known_poe_clients and (
+        poe_client_id = f"{POE_SWITCH}-{mac}"
-                mac in controller.wireless_clients
+        client = controller.api.clients[mac]
-                or client.sw_mac not in devices
+
-                or not devices[client.sw_mac].ports[client.sw_port].port_poe
+        if poe_client_id not in previously_known_poe_clients and (
-                or not devices[client.sw_mac].ports[client.sw_port].poe_enable
+            mac in controller.wireless_clients
-                or controller.mac == client.mac
+            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
+
+        # Multiple POE-devices on same port means non UniFi POE driven switch
+        multi_clients_on_port = False
+        for client2 in controller.api.clients.values():
+
+            if poe_client_id in previously_known_poe_clients:
+                break
+
+            if (
+                client2.is_wired
+                and client.mac != client2.mac
+                and client.sw_mac == client2.sw_mac
+                and client.sw_port == client2.sw_port
             ):
-                continue
+                multi_clients_on_port = True
+                break
 
-            # Multiple POE-devices on same port means non UniFi POE driven switch
+        if multi_clients_on_port:
-            multi_clients_on_port = False
+            continue
-            for client2 in controller.api.clients.values():
 
-                if poe_client_id in previously_known_poe_clients:
+        switches.append(UniFiPOEClientSwitch(client, controller))
-                    break
-
-                if (
-                    client2.is_wired
-                    and client.mac != client2.mac
-                    and client.sw_mac == client2.sw_mac
-                    and client.sw_port == client2.sw_port
-                ):
-                    multi_clients_on_port = True
-                    break
-
-            if multi_clients_on_port:
-                continue
-
-            switches.append(UniFiPOEClientSwitch(client, controller))
 
     if switches:
         async_add_entities(switches)

homeassistant/components/unifi/unifi_client.py

@@ -12,6 +12,7 @@ from aiounifi.events import (
     WIRELESS_CLIENT_CONNECTED,
     WIRELESS_CLIENT_DISCONNECTED,
     WIRELESS_CLIENT_ROAM,
+    WIRELESS_CLIENT_ROAMRADIO,
     WIRELESS_CLIENT_UNBLOCKED,
 )
 
@@ -25,7 +26,6 @@ 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_ROAMRADIO = "EVT_WU_RoamRadio"
 WIRELESS_CLIENT = (
     WIRELESS_CLIENT_CONNECTED,
     WIRELESS_CLIENT_DISCONNECTED,
@@ -55,7 +55,6 @@ class UniFiClient(UniFiBase):
     async def async_added_to_hass(self) -> None:
         """Client entity created."""
         await super().async_added_to_hass()
-        LOGGER.debug("New client %s (%s)", self.entity_id, self.client.mac)
         self.client.register_callback(self.async_update_callback)
 
     async def async_will_remove_from_hass(self) -> None:

homeassistant/components/unifi/unifi_entity_base.py

@@ -1,10 +1,13 @@
 """Base class for UniFi entities."""
+import logging
 
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_registry import async_entries_for_device
 
+LOGGER = logging.getLogger(__name__)
+
 
 class UniFiBase(Entity):
     """UniFi entity base class."""
@@ -27,6 +30,7 @@ class UniFiBase(Entity):
 
     async def async_added_to_hass(self) -> None:
         """Entity created."""
+        LOGGER.debug("New %s entity %s (%s)", self.TYPE, self.entity_id, self.mac)
         for signal, method in (
             (self.controller.signal_reachable, self.async_update_callback),
             (self.controller.signal_options_update, self.options_updated),

tests/components/unifi/test_device_tracker.py

@@ -2,7 +2,13 @@
 from copy import copy
 from datetime import timedelta
 
-from aiounifi.controller import MESSAGE_CLIENT_REMOVED, SIGNAL_CONNECTION_STATE
+from aiounifi.controller import (
+    MESSAGE_CLIENT,
+    MESSAGE_CLIENT_REMOVED,
+    MESSAGE_DEVICE,
+    MESSAGE_EVENT,
+    SIGNAL_CONNECTION_STATE,
+)
 from aiounifi.websocket import SIGNAL_DATA, STATE_DISCONNECTED, STATE_RUNNING
 
 from homeassistant import config_entries
@@ -24,8 +30,10 @@ 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 = {
+    "ap_mac": "00:00:00:00:02:01",
     "essid": "ssid",
     "hostname": "client_1",
     "ip": "10.0.0.1",
@@ -95,6 +103,38 @@ DEVICE_2 = {
     "version": "4.0.42.10433",
 }
 
+EVENT_CLIENT_1_WIRELESS_CONNECTED = {
+    "user": CLIENT_1["mac"],
+    "ssid": CLIENT_1["essid"],
+    "ap": CLIENT_1["ap_mac"],
+    "radio": "na",
+    "channel": "44",
+    "hostname": CLIENT_1["hostname"],
+    "key": "EVT_WU_Connected",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587753456179,
+    "datetime": "2020-04-24T18:37:36Z",
+    "msg": f'User{[CLIENT_1["mac"]]} has connected to AP[{CLIENT_1["ap_mac"]}] with SSID "{CLIENT_1["essid"]}" on "channel 44(na)"',
+    "_id": "5ea331fa30c49e00f90ddc1a",
+}
+
+EVENT_CLIENT_1_WIRELESS_DISCONNECTED = {
+    "user": CLIENT_1["mac"],
+    "ssid": CLIENT_1["essid"],
+    "hostname": CLIENT_1["hostname"],
+    "ap": CLIENT_1["ap_mac"],
+    "duration": 467,
+    "bytes": 459039,
+    "key": "EVT_WU_Disconnected",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587752927000,
+    "datetime": "2020-04-24T18:28:47Z",
+    "msg": f'User{[CLIENT_1["mac"]]} disconnected from "{CLIENT_1["essid"]}" (7m 47s connected, 448.28K bytes, last AP[{CLIENT_1["ap_mac"]}])',
+    "_id": "5ea32ff730c49e00f90dca1a",
+}
+
 
 async def test_platform_manually_configured(hass):
     """Test that nothing happens when configuring unifi through device tracker platform."""
@@ -114,6 +154,58 @@ async def test_no_clients(hass):
     assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 0
 
 
+async def test_tracked_wireless_clients(hass):
+    """Test the update_items function with some clients."""
+    controller = await setup_unifi_integration(hass, clients_response=[CLIENT_1])
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+    assert client_1.state == "not_home"
+
+    # 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],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
+    # State change signalling works with events
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_CLIENT_1_WIRELESS_DISCONNECTED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.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"
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_CLIENT_1_WIRELESS_CONNECTED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
+    # test wired bug
+
+
 async def test_tracked_devices(hass):
     """Test the update_items function with some clients."""
     client_4_copy = copy(CLIENT_4)
@@ -157,11 +249,11 @@ async def test_tracked_devices(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": "sta:sync"}, "data": [client_1_copy]}
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
     controller.api.message_handler(event)
     device_1_copy = copy(DEVICE_1)
     device_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
-    event = {"meta": {"message": "device:sync"}, "data": [device_1_copy]}
+    event = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_1_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
@@ -174,7 +266,7 @@ async def test_tracked_devices(hass):
     # Disabled device is unavailable
     device_1_copy = copy(DEVICE_1)
     device_1_copy["disabled"] = True
-    event = {"meta": {"message": "device:sync"}, "data": [device_1_copy]}
+    event = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_1_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
@@ -395,7 +487,7 @@ async def test_option_ssid_filter(hass):
 
     client_3_copy = copy(CLIENT_3)
     client_3_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
-    event = {"meta": {"message": "sta:sync"}, "data": [client_3_copy]}
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
@@ -407,7 +499,7 @@ async def test_option_ssid_filter(hass):
     hass.config_entries.async_update_entry(
         controller.config_entry, options={CONF_SSID_FILTER: []},
     )
-    event = {"meta": {"message": "sta:sync"}, "data": [client_3_copy]}
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_3_copy]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
@@ -433,7 +525,7 @@ async def test_wireless_client_go_wired_issue(hass):
 
     client_1_client["is_wired"] = True
     client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
-    event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 
@@ -444,7 +536,7 @@ async def test_wireless_client_go_wired_issue(hass):
     with patch.object(
         dt_util, "utcnow", return_value=(dt_util.utcnow() + timedelta(minutes=5)),
     ):
-        event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
+        event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_client]}
         controller.api.message_handler(event)
         await hass.async_block_till_done()
 

tests/components/unifi/test_sensor.py

@@ -1,7 +1,7 @@
 """UniFi sensor platform tests."""
 from copy import deepcopy
 
-from aiounifi.controller import MESSAGE_CLIENT_REMOVED
+from aiounifi.controller import MESSAGE_CLIENT, MESSAGE_CLIENT_REMOVED
 from aiounifi.websocket import SIGNAL_DATA
 
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
@@ -99,7 +99,7 @@ async def test_sensors(hass):
     clients[1]["rx_bytes"] = 2345000000
     clients[1]["tx_bytes"] = 6789000000
 
-    event = {"meta": {"message": "sta:sync"}, "data": clients}
+    event = {"meta": {"message": MESSAGE_CLIENT}, "data": clients}
     controller.api.message_handler(event)
     await hass.async_block_till_done()
 

tests/components/unifi/test_switch.py

@@ -1,7 +1,7 @@
-"""UniFi POE control platform tests."""
+"""UniFi switch platform tests."""
 from copy import deepcopy
 
-from aiounifi.controller import MESSAGE_CLIENT_REMOVED
+from aiounifi.controller import MESSAGE_CLIENT_REMOVED, MESSAGE_EVENT
 from aiounifi.websocket import SIGNAL_DATA
 
 from homeassistant import config_entries
@@ -197,6 +197,35 @@ UNBLOCKED = {
     "oui": "Producer",
 }
 
+EVENT_BLOCKED_CLIENT_CONNECTED = {
+    "user": BLOCKED["mac"],
+    "radio": "na",
+    "channel": "44",
+    "hostname": BLOCKED["hostname"],
+    "key": "EVT_WU_Connected",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587753456179,
+    "datetime": "2020-04-24T18:37:36Z",
+    "msg": f'User{[BLOCKED["mac"]]} has connected."',
+    "_id": "5ea331fa30c49e00f90ddc1a",
+}
+
+
+EVENT_CLIENT_2_CONNECTED = {
+    "user": CLIENT_2["mac"],
+    "radio": "na",
+    "channel": "44",
+    "hostname": CLIENT_2["hostname"],
+    "key": "EVT_WU_Connected",
+    "subsystem": "wlan",
+    "site_id": "name",
+    "time": 1587753456179,
+    "datetime": "2020-04-24T18:37:36Z",
+    "msg": f'User{[CLIENT_2["mac"]]} has connected."',
+    "_id": "5ea331fa30c49e00f90ddc1a",
+}
+
 
 async def test_platform_manually_configured(hass):
     """Test that we do not discover anything or try to set up a controller."""
@@ -360,7 +389,16 @@ async def test_new_client_discovered_on_block_control(hass):
         "meta": {"message": "sta:sync"},
         "data": [BLOCKED],
     }
-    controller.api.session_handler("data")
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_BLOCKED_CLIENT_CONNECTED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
     await hass.async_block_till_done()
 
     assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@@ -423,9 +461,22 @@ async def test_new_client_discovered_on_poe_control(hass):
         "meta": {"message": "sta:sync"},
         "data": [CLIENT_2],
     }
-    controller.api.session_handler("data")
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_EVENT},
+        "data": [EVENT_CLIENT_2_CONNECTED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+    switch_2 = hass.states.get("switch.poe_client_2")
+    assert switch_2 is not None
 
-    # Calling a service will trigger the updates to run
     await hass.services.async_call(
         SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.poe_client_1"}, blocking=True
     )
@@ -453,10 +504,6 @@ async def test_new_client_discovered_on_poe_control(hass):
         "path": "/rest/device/mock-id",
     }
 
-    switch_2 = hass.states.get("switch.poe_client_2")
-    assert switch_2 is not None
-    assert switch_2.state == "on"
-
 
 async def test_ignore_multiple_poe_clients_on_same_port(hass):
     """Ignore when there are multiple POE driven clients on same port.