diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 0877cda7475..2394dfe92d8 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/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() diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index 7f70d4c9f37..a32fc42715f 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/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" -} +} \ No newline at end of file diff --git a/homeassistant/components/unifi/services.py b/homeassistant/components/unifi/services.py new file mode 100644 index 00000000000..dca95a764c3 --- /dev/null +++ b/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) diff --git a/homeassistant/components/unifi/services.yaml b/homeassistant/components/unifi/services.yaml new file mode 100644 index 00000000000..435661afd4a --- /dev/null +++ b/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. diff --git a/requirements_all.txt b/requirements_all.txt index fa742db898f..def12433ec2 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index d5e02428222..3edea51b6ca 100644 --- a/requirements_test_all.txt +++ b/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 diff --git a/tests/components/unifi/test_services.py b/tests/components/unifi/test_services.py new file mode 100644 index 00000000000..388a33a4c64 --- /dev/null +++ b/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