From 8d6c4e33064d2b23eccd1b6c18dc621881720dfa Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 25 Jul 2023 14:01:57 +0200 Subject: [PATCH] Add controls to enable and disable a UniFi WLAN (#97204) --- homeassistant/components/unifi/entity.py | 20 ++++- homeassistant/components/unifi/image.py | 27 ++---- homeassistant/components/unifi/switch.py | 30 +++++++ tests/components/unifi/test_switch.py | 105 +++++++++++++++++++++++ 4 files changed, 161 insertions(+), 21 deletions(-) diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index 18a132be6a8..70b28e34dd0 100644 --- a/homeassistant/components/unifi/entity.py +++ b/homeassistant/components/unifi/entity.py @@ -18,11 +18,14 @@ from aiounifi.models.event import Event, EventKey from homeassistant.core import callback from homeassistant.helpers import entity_registry as er -from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC +from homeassistant.helpers.device_registry import ( + CONNECTION_NETWORK_MAC, + DeviceEntryType, +) from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity import DeviceInfo, Entity, EntityDescription -from .const import ATTR_MANUFACTURER +from .const import ATTR_MANUFACTURER, DOMAIN if TYPE_CHECKING: from .controller import UniFiController @@ -58,6 +61,19 @@ def async_device_device_info_fn(api: aiounifi.Controller, obj_id: str) -> Device ) +@callback +def async_wlan_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: + """Create device registry entry for WLAN.""" + wlan = api.wlans[obj_id] + return DeviceInfo( + entry_type=DeviceEntryType.SERVICE, + identifiers={(DOMAIN, wlan.id)}, + manufacturer=ATTR_MANUFACTURER, + model="UniFi WLAN", + name=wlan.name, + ) + + @dataclass class UnifiDescription(Generic[HandlerT, ApiItemT]): """Validate and load entities from different UniFi handlers.""" diff --git a/homeassistant/components/unifi/image.py b/homeassistant/components/unifi/image.py index 730720753d4..c26f06cb5f2 100644 --- a/homeassistant/components/unifi/image.py +++ b/homeassistant/components/unifi/image.py @@ -8,24 +8,26 @@ from collections.abc import Callable from dataclasses import dataclass from typing import Generic -import aiounifi from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.wlans import Wlans from aiounifi.models.api import ApiItemT from aiounifi.models.wlan import Wlan -from homeassistant.components.image import DOMAIN, ImageEntity, ImageEntityDescription +from homeassistant.components.image import ImageEntity, ImageEntityDescription from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback -from homeassistant.helpers.device_registry import DeviceEntryType -from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback import homeassistant.util.dt as dt_util -from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN +from .const import DOMAIN as UNIFI_DOMAIN from .controller import UniFiController -from .entity import HandlerT, UnifiEntity, UnifiEntityDescription +from .entity import ( + HandlerT, + UnifiEntity, + UnifiEntityDescription, + async_wlan_device_info_fn, +) @callback @@ -34,19 +36,6 @@ def async_wlan_qr_code_image_fn(controller: UniFiController, wlan: Wlan) -> byte return controller.api.wlans.generate_wlan_qr_code(wlan) -@callback -def async_wlan_device_info_fn(api: aiounifi.Controller, obj_id: str) -> DeviceInfo: - """Create device registry entry for WLAN.""" - wlan = api.wlans[obj_id] - return DeviceInfo( - entry_type=DeviceEntryType.SERVICE, - identifiers={(DOMAIN, wlan.id)}, - manufacturer=ATTR_MANUFACTURER, - model="UniFi Network", - name=wlan.name, - ) - - @dataclass class UnifiImageEntityDescriptionMixin(Generic[HandlerT, ApiItemT]): """Validate and load entities from different UniFi handlers.""" diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 846c6d12234..ca11cdfea30 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -3,6 +3,7 @@ Support for controlling power supply of clients which are powered over Ethernet (POE). Support for controlling network access of clients selected in option flow. Support for controlling deep packet inspection (DPI) restriction groups. +Support for controlling WLAN availability. """ from __future__ import annotations @@ -17,6 +18,7 @@ from aiounifi.interfaces.clients import Clients from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups from aiounifi.interfaces.outlets import Outlets from aiounifi.interfaces.ports import Ports +from aiounifi.interfaces.wlans import Wlans from aiounifi.models.api import ApiItemT from aiounifi.models.client import Client, ClientBlockRequest from aiounifi.models.device import ( @@ -28,6 +30,7 @@ from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup from aiounifi.models.event import Event, EventKey from aiounifi.models.outlet import Outlet from aiounifi.models.port import Port +from aiounifi.models.wlan import Wlan, WlanEnableRequest from homeassistant.components.switch import ( DOMAIN, @@ -54,6 +57,7 @@ from .entity import ( UnifiEntityDescription, async_device_available_fn, async_device_device_info_fn, + async_wlan_device_info_fn, ) CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED) @@ -137,6 +141,13 @@ async def async_poe_port_control_fn( await api.request(DeviceSetPoePortModeRequest.create(device, int(index), state)) +async def async_wlan_control_fn( + api: aiounifi.Controller, obj_id: str, target: bool +) -> None: + """Control outlet relay.""" + await api.request(WlanEnableRequest.create(obj_id, target)) + + @dataclass class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, ApiItemT]): """Validate and load entities from different UniFi handlers.""" @@ -233,6 +244,25 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe, unique_id_fn=lambda controller, obj_id: f"{obj_id.split('_', 1)[0]}-poe-{obj_id.split('_', 1)[1]}", ), + UnifiSwitchEntityDescription[Wlans, Wlan]( + key="WLAN control", + device_class=SwitchDeviceClass.SWITCH, + entity_category=EntityCategory.CONFIG, + has_entity_name=True, + icon="mdi:wifi-check", + allowed_fn=lambda controller, obj_id: True, + api_handler_fn=lambda api: api.wlans, + available_fn=lambda controller, _: controller.available, + control_fn=async_wlan_control_fn, + device_info_fn=async_wlan_device_info_fn, + event_is_on=None, + event_to_subscribe=None, + is_on_fn=lambda controller, wlan: wlan.enabled, + name_fn=lambda wlan: None, + object_fn=lambda api, obj_id: api.wlans[obj_id], + supported_fn=lambda controller, obj_id: True, + unique_id_fn=lambda controller, obj_id: f"wlan-{obj_id}", + ), ) diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index f93abc291b8..ad5131614af 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -580,6 +580,43 @@ OUTLET_UP1 = { } +WLAN = { + "_id": "012345678910111213141516", + "bc_filter_enabled": False, + "bc_filter_list": [], + "dtim_mode": "default", + "dtim_na": 1, + "dtim_ng": 1, + "enabled": True, + "group_rekey": 3600, + "mac_filter_enabled": False, + "mac_filter_list": [], + "mac_filter_policy": "allow", + "minrate_na_advertising_rates": False, + "minrate_na_beacon_rate_kbps": 6000, + "minrate_na_data_rate_kbps": 6000, + "minrate_na_enabled": False, + "minrate_na_mgmt_rate_kbps": 6000, + "minrate_ng_advertising_rates": False, + "minrate_ng_beacon_rate_kbps": 1000, + "minrate_ng_data_rate_kbps": 1000, + "minrate_ng_enabled": False, + "minrate_ng_mgmt_rate_kbps": 1000, + "name": "SSID 1", + "no2ghz_oui": False, + "schedule": [], + "security": "wpapsk", + "site_id": "5a32aa4ee4b0412345678910", + "usergroup_id": "012345678910111213141518", + "wep_idx": 1, + "wlangroup_id": "012345678910111213141519", + "wpa_enc": "ccmp", + "wpa_mode": "wpa2", + "x_iapp_key": "01234567891011121314151617181920", + "x_passphrase": "password", +} + + async def test_no_clients( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker ) -> None: @@ -1230,3 +1267,71 @@ async def test_remove_poe_client_switches( for entry in ent_reg.entities.values() if entry.config_entry_id == config_entry.entry_id ] + + +async def test_wlan_switches( + hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket +) -> None: + """Test control of UniFi WLAN availability.""" + config_entry = await setup_unifi_integration( + hass, aioclient_mock, wlans_response=[WLAN] + ) + controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 + + ent_reg = er.async_get(hass) + ent_reg_entry = ent_reg.async_get("switch.ssid_1") + assert ent_reg_entry.unique_id == "wlan-012345678910111213141516" + assert ent_reg_entry.entity_category is EntityCategory.CONFIG + + # Validate state object + switch_1 = hass.states.get("switch.ssid_1") + assert switch_1 is not None + assert switch_1.state == STATE_ON + assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.SWITCH + + # Update state object + wlan = deepcopy(WLAN) + wlan["enabled"] = False + mock_unifi_websocket(message=MessageKey.WLAN_CONF_UPDATED, data=wlan) + await hass.async_block_till_done() + assert hass.states.get("switch.ssid_1").state == STATE_OFF + + # Disable WLAN + aioclient_mock.clear_requests() + aioclient_mock.put( + f"https://{controller.host}:1234/api/s/{controller.site}" + + f"/rest/wlanconf/{WLAN['_id']}", + ) + + await hass.services.async_call( + SWITCH_DOMAIN, + "turn_off", + {"entity_id": "switch.ssid_1"}, + blocking=True, + ) + assert aioclient_mock.call_count == 1 + assert aioclient_mock.mock_calls[0][2] == {"enabled": False} + + # Enable WLAN + await hass.services.async_call( + SWITCH_DOMAIN, + "turn_on", + {"entity_id": "switch.ssid_1"}, + blocking=True, + ) + assert aioclient_mock.call_count == 2 + assert aioclient_mock.mock_calls[1][2] == {"enabled": True} + + # Availability signalling + + # Controller disconnects + mock_unifi_websocket(state=WebsocketState.DISCONNECTED) + await hass.async_block_till_done() + assert hass.states.get("switch.ssid_1").state == STATE_UNAVAILABLE + + # Controller reconnects + mock_unifi_websocket(state=WebsocketState.RUNNING) + await hass.async_block_till_done() + assert hass.states.get("switch.ssid_1").state == STATE_OFF