Service to remove clients from UniFi Controller (#56717)
This commit is contained in:
parent
e729339538
commit
d61a9e8b72
7 changed files with 238 additions and 5 deletions
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
@ -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"
|
||||
}
|
||||
}
|
69
homeassistant/components/unifi/services.py
Normal file
69
homeassistant/components/unifi/services.py
Normal file
|
@ -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)
|
3
homeassistant/components/unifi/services.yaml
Normal file
3
homeassistant/components/unifi/services.yaml
Normal file
|
@ -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.
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
151
tests/components/unifi/test_services.py
Normal file
151
tests/components/unifi/test_services.py
Normal file
|
@ -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
|
Loading…
Add table
Reference in a new issue