Service to remove clients from UniFi Controller (#56717)

homeassistant/components/unifi/__init__.py

@@ -11,6 +11,7 @@ from .const import (
     UNIFI_WIRELESS_CLIENTS,
 )
 from .controller import UniFiController
+from .services import async_setup_services, async_unload_services
 
 SAVE_DELAY = 10
 STORAGE_KEY = "unifi_data"
@@ -43,6 +44,7 @@ async def async_setup_entry(hass, config_entry):
         )
 
     hass.data[UNIFI_DOMAIN][config_entry.entry_id] = controller
+    await async_setup_services(hass)
 
     config_entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)
@@ -68,6 +70,10 @@ async def async_setup_entry(hass, config_entry):
 async def async_unload_entry(hass, config_entry):
     """Unload a config entry."""
     controller = hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id)
+
+    if not hass.data[UNIFI_DOMAIN]:
+        await async_unload_services(hass)
+
     return await controller.async_reset()
 
 

homeassistant/components/unifi/manifest.json

@@ -3,8 +3,12 @@
   "name": "Ubiquiti UniFi",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/unifi",
-  "requirements": ["aiounifi==26"],
-  "codeowners": ["@Kane610"],
+  "requirements": [
+    "aiounifi==27"
+  ],
+  "codeowners": [
+    "@Kane610"
+  ],
   "quality_scale": "platinum",
   "ssdp": [
     {
@@ -19,4 +23,4 @@
     }
   ],
   "iot_class": "local_push"
-}
+}

homeassistant/components/unifi/services.py

@@ -0,0 +1,69 @@
+"""UniFi services."""
+
+from .const import DOMAIN as UNIFI_DOMAIN
+
+UNIFI_SERVICES = "unifi_services"
+
+SERVICE_REMOVE_CLIENTS = "remove_clients"
+
+
+async def async_setup_services(hass) -> None:
+    """Set up services for UniFi integration."""
+    if hass.data.get(UNIFI_SERVICES, False):
+        return
+
+    hass.data[UNIFI_SERVICES] = True
+
+    async def async_call_unifi_service(service_call) -> None:
+        """Call correct UniFi service."""
+        service = service_call.service
+        service_data = service_call.data
+
+        controllers = hass.data[UNIFI_DOMAIN].values()
+
+        if service == SERVICE_REMOVE_CLIENTS:
+            await async_remove_clients(controllers, service_data)
+
+    hass.services.async_register(
+        UNIFI_DOMAIN,
+        SERVICE_REMOVE_CLIENTS,
+        async_call_unifi_service,
+    )
+
+
+async def async_unload_services(hass) -> None:
+    """Unload UniFi services."""
+    if not hass.data.get(UNIFI_SERVICES):
+        return
+
+    hass.data[UNIFI_SERVICES] = False
+
+    hass.services.async_remove(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS)
+
+
+async def async_remove_clients(controllers, data) -> None:
+    """Remove select clients from controller.
+
+    Validates based on:
+    - Total time between first seen and last seen is less than 15 minutes.
+    - Neither IP, hostname nor name is configured.
+    """
+    for controller in controllers:
+
+        if not controller.available:
+            continue
+
+        clients_to_remove = []
+
+        for client in controller.api.clients_all.values():
+
+            if client.last_seen - client.first_seen > 900:
+                continue
+
+            if any({client.fixed_ip, client.hostname, client.name}):
+                continue
+
+            clients_to_remove.append(client.mac)
+
+        if clients_to_remove:
+            await controller.api.clients.remove_clients(macs=clients_to_remove)

homeassistant/components/unifi/services.yaml

@@ -0,0 +1,3 @@
+remove_clients:
+  name: Remove clients from the UniFi Controller
+  description: Clean up clients that has only been associated with the controller for a short period of time.

requirements_all.txt

@@ -255,7 +255,7 @@ aiosyncthing==0.5.1
 aiotractive==0.5.2
 
 # homeassistant.components.unifi
-aiounifi==26
+aiounifi==27
 
 # homeassistant.components.watttime
 aiowatttime==0.1.1

requirements_test_all.txt

@@ -179,7 +179,7 @@ aiosyncthing==0.5.1
 aiotractive==0.5.2
 
 # homeassistant.components.unifi
-aiounifi==26
+aiounifi==27
 
 # homeassistant.components.watttime
 aiowatttime==0.1.1

tests/components/unifi/test_services.py

@@ -0,0 +1,151 @@
+"""deCONZ service tests."""
+
+from unittest.mock import Mock, patch
+
+from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
+from homeassistant.components.unifi.services import (
+    SERVICE_REMOVE_CLIENTS,
+    UNIFI_SERVICES,
+    async_setup_services,
+    async_unload_services,
+)
+
+from .test_controller import setup_unifi_integration
+
+
+async def test_service_setup(hass):
+    """Verify service setup works."""
+    assert UNIFI_SERVICES not in hass.data
+    with patch(
+        "homeassistant.core.ServiceRegistry.async_register", return_value=Mock(True)
+    ) as async_register:
+        await async_setup_services(hass)
+        assert hass.data[UNIFI_SERVICES] is True
+        assert async_register.call_count == 1
+
+
+async def test_service_setup_already_registered(hass):
+    """Make sure that services are only registered once."""
+    hass.data[UNIFI_SERVICES] = True
+    with patch(
+        "homeassistant.core.ServiceRegistry.async_register", return_value=Mock(True)
+    ) as async_register:
+        await async_setup_services(hass)
+        async_register.assert_not_called()
+
+
+async def test_service_unload(hass):
+    """Verify service unload works."""
+    hass.data[UNIFI_SERVICES] = True
+    with patch(
+        "homeassistant.core.ServiceRegistry.async_remove", return_value=Mock(True)
+    ) as async_remove:
+        await async_unload_services(hass)
+        assert hass.data[UNIFI_SERVICES] is False
+        assert async_remove.call_count == 1
+
+
+async def test_service_unload_not_registered(hass):
+    """Make sure that services can only be unloaded once."""
+    with patch(
+        "homeassistant.core.ServiceRegistry.async_remove", return_value=Mock(True)
+    ) as async_remove:
+        await async_unload_services(hass)
+        assert UNIFI_SERVICES not in hass.data
+        async_remove.assert_not_called()
+
+
+async def test_remove_clients(hass, aioclient_mock):
+    """Verify removing different variations of clients work."""
+    clients = [
+        {
+            "first_seen": 100,
+            "last_seen": 500,
+            "mac": "00:00:00:00:00:01",
+        },
+        {
+            "first_seen": 100,
+            "last_seen": 1100,
+            "mac": "00:00:00:00:00:02",
+        },
+        {
+            "first_seen": 100,
+            "last_seen": 500,
+            "fixed_ip": "1.2.3.4",
+            "mac": "00:00:00:00:00:03",
+        },
+        {
+            "first_seen": 100,
+            "last_seen": 500,
+            "hostname": "hostname",
+            "mac": "00:00:00:00:00:04",
+        },
+        {
+            "first_seen": 100,
+            "last_seen": 500,
+            "name": "name",
+            "mac": "00:00:00:00:00:05",
+        },
+    ]
+    config_entry = await setup_unifi_integration(
+        hass, aioclient_mock, clients_all_response=clients
+    )
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
+    aioclient_mock.clear_requests()
+    aioclient_mock.post(
+        f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr",
+    )
+
+    await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True)
+    assert aioclient_mock.mock_calls[0][2] == {
+        "cmd": "forget-sta",
+        "macs": ["00:00:00:00:00:01"],
+    }
+
+
+async def test_remove_clients_controller_unavailable(hass, aioclient_mock):
+    """Verify no call is made if controller is unavailable."""
+    clients = [
+        {
+            "first_seen": 100,
+            "last_seen": 500,
+            "mac": "00:00:00:00:00:01",
+        }
+    ]
+    config_entry = await setup_unifi_integration(
+        hass, aioclient_mock, clients_all_response=clients
+    )
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+    controller.available = False
+
+    aioclient_mock.clear_requests()
+    aioclient_mock.post(
+        f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr",
+    )
+
+    await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True)
+    assert aioclient_mock.call_count == 0
+
+
+async def test_remove_clients_no_call_on_empty_list(hass, aioclient_mock):
+    """Verify no call is made if no fitting client has been added to the list."""
+    clients = [
+        {
+            "first_seen": 100,
+            "last_seen": 1100,
+            "mac": "00:00:00:00:00:01",
+        }
+    ]
+    config_entry = await setup_unifi_integration(
+        hass, aioclient_mock, clients_all_response=clients
+    )
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
+    aioclient_mock.clear_requests()
+    aioclient_mock.post(
+        f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr",
+    )
+
+    await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True)
+    assert aioclient_mock.call_count == 0