UniFi - Support automatic removal of clients (#34307)

homeassistant/components/unifi/controller.py

@@ -5,7 +5,14 @@ import ssl
 
 from aiohttp import CookieJar
 import aiounifi
-from aiounifi.controller import SIGNAL_CONNECTION_STATE
+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.websocket import STATE_DISCONNECTED, STATE_RUNNING
 import async_timeout
@@ -161,16 +168,23 @@ class UniFiController:
                 if not self.available:
                     self.hass.loop.call_later(RETRY_TIMER, self.reconnect)
 
-        elif signal == "new_data" and data:
+        elif signal == SIGNAL_DATA and data:
-            if "event" in data:
+
-                if data["event"].event in (
+            if DATA_EVENT in data:
+                if data[DATA_EVENT].event in (
                     WIRELESS_CLIENT_CONNECTED,
                     WIRELESS_GUEST_CONNECTED,
                 ):
                     self.update_wireless_clients()
-            elif "clients" in data or "devices" in data:
+
+            elif DATA_CLIENT in data or DATA_DEVICE in data:
                 async_dispatcher_send(self.hass, self.signal_update)
 
+            elif DATA_CLIENT_REMOVED in data:
+                async_dispatcher_send(
+                    self.hass, self.signal_remove, data[DATA_CLIENT_REMOVED]
+                )
+
     @property
     def signal_reachable(self) -> str:
         """Integration specific event to signal a change in connection status."""
@@ -181,6 +195,11 @@ class UniFiController:
         """Event specific per UniFi entry to signal new data."""
         return f"unifi-update-{self.controller_id}"
 
+    @property
+    def signal_remove(self):
+        """Event specific per UniFi entry to signal removal of entities."""
+        return f"unifi-remove-{self.controller_id}"
+
     @property
     def signal_options_update(self):
         """Event specific per UniFi entry to signal new options."""

homeassistant/components/unifi/device_tracker.py

@@ -49,10 +49,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     option_track_wired_clients = controller.option_track_wired_clients
     option_ssid_filter = controller.option_ssid_filter
 
-    registry = await hass.helpers.entity_registry.async_get_registry()
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
 
     # Restore clients that is not a part of active clients list.
-    for entity in registry.entities.values():
+    for entity in entity_registry.entities.values():
 
         if (
             entity.config_entry_id == config_entry.entry_id
@@ -69,7 +69,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
             controller.api.clients.process_raw([client.raw])
 
     @callback
-    def update_controller():
+    def items_added():
         """Update the values of the controller."""
         nonlocal option_track_clients
         nonlocal option_track_devices
@@ -80,7 +80,16 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         add_entities(controller, async_add_entities, tracked)
 
     controller.listeners.append(
-        async_dispatcher_connect(hass, controller.signal_update, update_controller)
+        async_dispatcher_connect(hass, controller.signal_update, items_added)
+    )
+
+    @callback
+    def items_removed(mac_addresses: set) -> None:
+        """Items have been removed from the controller."""
+        remove_entities(controller, mac_addresses, tracked, entity_registry)
+
+    controller.listeners.append(
+        async_dispatcher_connect(hass, controller.signal_remove, items_removed)
     )
 
     @callback
@@ -136,16 +145,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         option_track_devices = controller.option_track_devices
         option_track_wired_clients = controller.option_track_wired_clients
 
-        for mac in remove:
+        remove_entities(controller, remove, tracked, entity_registry)
-            entity = tracked.pop(mac)
-
-            if registry.async_is_registered(entity.entity_id):
-                registry.async_remove(entity.entity_id)
-
-            hass.async_create_task(entity.async_remove())
 
         if update:
-            update_controller()
+            items_added()
 
     controller.listeners.append(
         async_dispatcher_connect(
@@ -153,7 +156,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         )
     )
 
-    update_controller()
+    items_added()
 
 
 @callback
@@ -193,6 +196,18 @@ def add_entities(controller, async_add_entities, tracked):
         async_add_entities(new_tracked)
 
 
+@callback
+def remove_entities(controller, mac_addresses, tracked, entity_registry):
+    """Remove select tracked entities."""
+    for mac in mac_addresses:
+
+        if mac not in tracked:
+            continue
+
+        entity = tracked.pop(mac)
+        controller.hass.async_create_task(entity.async_remove())
+
+
 class UniFiClientTracker(UniFiClient, ScannerEntity):
     """Representation of a network client."""
 

homeassistant/components/unifi/manifest.json

@@ -3,11 +3,7 @@
   "name": "Ubiquiti UniFi",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/unifi",
-  "requirements": [
+  "requirements": ["aiounifi==18"],
-    "aiounifi==17"
+  "codeowners": ["@kane610"],
-  ],
-  "codeowners": [
-    "@kane610"
-  ],
   "quality_scale": "platinum"
 }

homeassistant/components/unifi/sensor.py

@@ -25,7 +25,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     entity_registry = await hass.helpers.entity_registry.async_get_registry()
 
     @callback
-    def update_controller():
+    def items_added():
         """Update the values of the controller."""
         nonlocal option_allow_bandwidth_sensors
 
@@ -35,7 +35,16 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         add_entities(controller, async_add_entities, sensors)
 
     controller.listeners.append(
-        async_dispatcher_connect(hass, controller.signal_update, update_controller)
+        async_dispatcher_connect(hass, controller.signal_update, items_added)
+    )
+
+    @callback
+    def items_removed(mac_addresses: set) -> None:
+        """Items have been removed from the controller."""
+        remove_entities(controller, mac_addresses, sensors, entity_registry)
+
+    controller.listeners.append(
+        async_dispatcher_connect(hass, controller.signal_remove, items_removed)
     )
 
     @callback
@@ -47,14 +56,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
             option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors
 
             if option_allow_bandwidth_sensors:
-                update_controller()
+                items_added()
 
             else:
                 for sensor in sensors.values():
-
-                    if entity_registry.async_is_registered(sensor.entity_id):
-                        entity_registry.async_remove(sensor.entity_id)
-
                     hass.async_create_task(sensor.async_remove())
 
                 sensors.clear()
@@ -65,7 +70,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         )
     )
 
-    update_controller()
+    items_added()
 
 
 @callback
@@ -92,6 +97,21 @@ def add_entities(controller, async_add_entities, sensors):
         async_add_entities(new_sensors)
 
 
+@callback
+def remove_entities(controller, mac_addresses, sensors, entity_registry):
+    """Remove select sensor entities."""
+    for mac in mac_addresses:
+
+        for direction in ("rx", "tx"):
+            item_id = f"{direction}-{mac}"
+
+            if item_id not in sensors:
+                continue
+
+            entity = sensors.pop(item_id)
+            controller.hass.async_create_task(entity.async_remove())
+
+
 class UniFiRxBandwidthSensor(UniFiClient):
     """Receiving bandwidth sensor."""
 

homeassistant/components/unifi/switch.py

@@ -55,12 +55,21 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
                 continue
 
     @callback
-    def update_controller():
+    def items_added():
         """Update the values of the controller."""
         add_entities(controller, async_add_entities, switches, switches_off)
 
     controller.listeners.append(
-        async_dispatcher_connect(hass, controller.signal_update, update_controller)
+        async_dispatcher_connect(hass, controller.signal_update, items_added)
+    )
+
+    @callback
+    def items_removed(mac_addresses: set) -> None:
+        """Items have been removed from the controller."""
+        remove_entities(controller, mac_addresses, switches, entity_registry)
+
+    controller.listeners.append(
+        async_dispatcher_connect(hass, controller.signal_remove, items_removed)
     )
 
     @callback
@@ -96,14 +105,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
 
         for client_id in remove:
             entity = switches.pop(client_id)
-
-            if entity_registry.async_is_registered(entity.entity_id):
-                entity_registry.async_remove(entity.entity_id)
-
             hass.async_create_task(entity.async_remove())
 
         if len(update) != len(option_block_clients):
-            update_controller()
+            items_added()
 
     controller.listeners.append(
         async_dispatcher_connect(
@@ -111,7 +116,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         )
     )
 
-    update_controller()
+    items_added()
     switches_off.clear()
 
 
@@ -189,6 +194,21 @@ def add_entities(controller, async_add_entities, switches, switches_off):
         async_add_entities(new_switches)
 
 
+@callback
+def remove_entities(controller, mac_addresses, switches, entity_registry):
+    """Remove select switch entities."""
+    for mac in mac_addresses:
+
+        for switch_type in ("block", "poe"):
+            item_id = f"{switch_type}-{mac}"
+
+            if item_id not in switches:
+                continue
+
+            entity = switches.pop(item_id)
+            controller.hass.async_create_task(entity.async_remove())
+
+
 class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
     """Representation of a client that uses POE."""
 

requirements_all.txt

@@ -211,7 +211,7 @@ aiopylgtv==0.3.3
 aioswitcher==1.1.0
 
 # homeassistant.components.unifi
-aiounifi==17
+aiounifi==18
 
 # homeassistant.components.wwlln
 aiowwlln==2.0.2

requirements_test_all.txt

@@ -94,7 +94,7 @@ aiopylgtv==0.3.3
 aioswitcher==1.1.0
 
 # homeassistant.components.unifi
-aiounifi==17
+aiounifi==18
 
 # homeassistant.components.wwlln
 aiowwlln==2.0.2

tests/components/unifi/test_controller.py

@@ -177,6 +177,7 @@ async def test_controller_setup(hass):
     assert controller.mac is None
 
     assert controller.signal_update == "unifi-update-1.2.3.4-site_id"
+    assert controller.signal_remove == "unifi-remove-1.2.3.4-site_id"
     assert controller.signal_options_update == "unifi-options-1.2.3.4-site_id"
 
 
@@ -206,7 +207,7 @@ async def test_reset_after_successful_setup(hass):
     """Calling reset when the entry has been setup."""
     controller = await setup_unifi_integration(hass)
 
-    assert len(controller.listeners) == 6
+    assert len(controller.listeners) == 9
 
     result = await controller.async_reset()
     await hass.async_block_till_done()

tests/components/unifi/test_device_tracker.py

@@ -2,8 +2,8 @@
 from copy import copy
 from datetime import timedelta
 
-from aiounifi.controller import SIGNAL_CONNECTION_STATE
+from aiounifi.controller import MESSAGE_CLIENT_REMOVED, SIGNAL_CONNECTION_STATE
-from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
+from aiounifi.websocket import SIGNAL_DATA, STATE_DISCONNECTED, STATE_RUNNING
 from asynctest import patch
 
 from homeassistant import config_entries
@@ -180,6 +180,35 @@ async def test_tracked_devices(hass):
     assert device_1.state == STATE_UNAVAILABLE
 
 
+async def test_remove_clients(hass):
+    """Test the remove_items function with some clients."""
+    controller = await setup_unifi_integration(
+        hass, clients_response=[CLIENT_1, CLIENT_2]
+    )
+    assert len(hass.states.async_entity_ids("device_tracker")) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    wired_client = hass.states.get("device_tracker.wired_client")
+    assert wired_client is not None
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_CLIENT_REMOVED},
+        "data": [CLIENT_1],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids("device_tracker")) == 1
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is None
+
+    wired_client = hass.states.get("device_tracker.wired_client")
+    assert wired_client is not None
+
+
 async def test_controller_state_change(hass):
     """Verify entities state reflect on controller becoming unavailable."""
     controller = await setup_unifi_integration(

tests/components/unifi/test_sensor.py

@@ -1,6 +1,9 @@
 """UniFi sensor platform tests."""
 from copy import deepcopy
 
+from aiounifi.controller import MESSAGE_CLIENT_REMOVED
+from aiounifi.websocket import SIGNAL_DATA
+
 from homeassistant.components import unifi
 import homeassistant.components.sensor as sensor
 from homeassistant.setup import async_setup_component
@@ -123,3 +126,44 @@ async def test_sensors(hass):
 
     wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
     assert wireless_client_tx.state == "6789.0"
+
+
+async def test_remove_sensors(hass):
+    """Test the remove_items function with some clients."""
+    controller = await setup_unifi_integration(
+        hass,
+        options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
+        clients_response=CLIENTS,
+    )
+    assert len(hass.states.async_entity_ids("sensor")) == 4
+    assert len(hass.states.async_entity_ids("device_tracker")) == 2
+
+    wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
+    assert wired_client_rx is not None
+    wired_client_tx = hass.states.get("sensor.wired_client_name_tx")
+    assert wired_client_tx is not None
+
+    wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
+    assert wireless_client_rx is not None
+    wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
+    assert wireless_client_tx is not None
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_CLIENT_REMOVED},
+        "data": [CLIENTS[0]],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids("sensor")) == 2
+    assert len(hass.states.async_entity_ids("device_tracker")) == 1
+
+    wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
+    assert wired_client_rx is None
+    wired_client_tx = hass.states.get("sensor.wired_client_name_tx")
+    assert wired_client_tx is None
+
+    wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
+    assert wireless_client_rx is not None
+    wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
+    assert wireless_client_tx is not None

tests/components/unifi/test_switch.py

@@ -1,6 +1,9 @@
 """UniFi POE control platform tests."""
 from copy import deepcopy
 
+from aiounifi.controller import MESSAGE_CLIENT_REMOVED
+from aiounifi.websocket import SIGNAL_DATA
+
 from homeassistant import config_entries
 from homeassistant.components import unifi
 import homeassistant.components.switch as switch
@@ -173,6 +176,7 @@ BLOCKED = {
     "ip": "10.0.0.1",
     "is_guest": False,
     "is_wired": False,
+    "last_seen": 1562600145,
     "mac": "00:00:00:00:01:01",
     "name": "Block Client 1",
     "noted": True,
@@ -184,6 +188,7 @@ UNBLOCKED = {
     "ip": "10.0.0.2",
     "is_guest": False,
     "is_wired": True,
+    "last_seen": 1562600145,
     "mac": "00:00:00:00:01:02",
     "name": "Block Client 2",
     "noted": True,
@@ -300,6 +305,38 @@ async def test_switches(hass):
     }
 
 
+async def test_remove_switches(hass):
+    """Test the update_items function with some clients."""
+    controller = await setup_unifi_integration(
+        hass,
+        options={CONF_BLOCK_CLIENT: [UNBLOCKED["mac"]]},
+        clients_response=[CLIENT_1, UNBLOCKED],
+        devices_response=[DEVICE_1],
+    )
+    assert len(hass.states.async_entity_ids("switch")) == 2
+
+    poe_switch = hass.states.get("switch.poe_client_1")
+    assert poe_switch is not None
+
+    block_switch = hass.states.get("switch.block_client_2")
+    assert block_switch is not None
+
+    controller.api.websocket._data = {
+        "meta": {"message": MESSAGE_CLIENT_REMOVED},
+        "data": [CLIENT_1, UNBLOCKED],
+    }
+    controller.api.session_handler(SIGNAL_DATA)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids("switch")) == 0
+
+    poe_switch = hass.states.get("switch.poe_client_1")
+    assert poe_switch is None
+
+    block_switch = hass.states.get("switch.block_client_2")
+    assert block_switch is None
+
+
 async def test_new_client_discovered_on_block_control(hass):
     """Test if 2nd update has a new client."""
     controller = await setup_unifi_integration(