diff --git a/homeassistant/components/unifi/button.py b/homeassistant/components/unifi/button.py index 45fc76c73df..86c38a5bf3d 100644 --- a/homeassistant/components/unifi/button.py +++ b/homeassistant/components/unifi/button.py @@ -7,12 +7,14 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from dataclasses import dataclass +import secrets from typing import Any import aiounifi from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.devices import Devices from aiounifi.interfaces.ports import Ports +from aiounifi.interfaces.wlans import Wlans from aiounifi.models.api import ApiItemT from aiounifi.models.device import ( Device, @@ -20,6 +22,7 @@ from aiounifi.models.device import ( DeviceRestartRequest, ) from aiounifi.models.port import Port +from aiounifi.models.wlan import Wlan, WlanChangePasswordRequest from homeassistant.components.button import ( ButtonDeviceClass, @@ -37,6 +40,8 @@ from .entity import ( UnifiEntityDescription, async_device_available_fn, async_device_device_info_fn, + async_wlan_available_fn, + async_wlan_device_info_fn, ) from .hub import UnifiHub @@ -56,6 +61,15 @@ async def async_power_cycle_port_control_fn( await api.request(DevicePowerCyclePortRequest.create(mac, int(index))) +async def async_regenerate_password_control_fn( + api: aiounifi.Controller, obj_id: str +) -> None: + """Regenerate WLAN password.""" + await api.request( + WlanChangePasswordRequest.create(obj_id, secrets.token_urlsafe(15)) + ) + + @dataclass(frozen=True, kw_only=True) class UnifiButtonEntityDescription( ButtonEntityDescription, UnifiEntityDescription[HandlerT, ApiItemT] @@ -91,6 +105,19 @@ ENTITY_DESCRIPTIONS: tuple[UnifiButtonEntityDescription, ...] = ( supported_fn=lambda hub, obj_id: bool(hub.api.ports[obj_id].port_poe), unique_id_fn=lambda hub, obj_id: f"power_cycle-{obj_id}", ), + UnifiButtonEntityDescription[Wlans, Wlan]( + key="WLAN regenerate password", + device_class=ButtonDeviceClass.UPDATE, + entity_category=EntityCategory.CONFIG, + entity_registry_enabled_default=False, + api_handler_fn=lambda api: api.wlans, + available_fn=async_wlan_available_fn, + control_fn=async_regenerate_password_control_fn, + device_info_fn=async_wlan_device_info_fn, + name_fn=lambda wlan: "Regenerate Password", + object_fn=lambda api, obj_id: api.wlans[obj_id], + unique_id_fn=lambda hub, obj_id: f"regenerate_password-{obj_id}", + ), ) @@ -109,7 +136,7 @@ async def async_setup_entry( class UnifiButtonEntity(UnifiEntity[HandlerT, ApiItemT], ButtonEntity): - """Base representation of a UniFi image.""" + """Base representation of a UniFi button.""" entity_description: UnifiButtonEntityDescription[HandlerT, ApiItemT] diff --git a/tests/components/unifi/test_button.py b/tests/components/unifi/test_button.py index d5be861139b..8f9838e3e37 100644 --- a/tests/components/unifi/test_button.py +++ b/tests/components/unifi/test_button.py @@ -1,20 +1,64 @@ """UniFi Network button platform tests.""" +from datetime import timedelta + from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, ButtonDeviceClass from homeassistant.components.unifi.const import CONF_SITE_ID +from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY from homeassistant.const import ( ATTR_DEVICE_CLASS, CONF_HOST, + CONTENT_TYPE_JSON, STATE_UNAVAILABLE, EntityCategory, ) from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er +from homeassistant.helpers.entity_registry import RegistryEntryDisabler +import homeassistant.util.dt as dt_util from .test_hub import setup_unifi_integration +from tests.common import async_fire_time_changed from tests.test_util.aiohttp import AiohttpClientMocker +WLAN_ID = "_id" +WLAN = { + 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_restart_device_button( hass: HomeAssistant, @@ -168,3 +212,71 @@ async def test_power_cycle_poe( assert ( hass.states.get("button.switch_port_1_power_cycle").state != STATE_UNAVAILABLE ) + + +async def test_wlan_regenerate_password( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + aioclient_mock: AiohttpClientMocker, + websocket_mock, +) -> None: + """Test WLAN regenerate password button.""" + + config_entry = await setup_unifi_integration( + hass, aioclient_mock, wlans_response=[WLAN] + ) + assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 0 + + button_regenerate_password = "button.ssid_1_regenerate_password" + + ent_reg_entry = entity_registry.async_get(button_regenerate_password) + assert ent_reg_entry.unique_id == "regenerate_password-012345678910111213141516" + assert ent_reg_entry.disabled_by == RegistryEntryDisabler.INTEGRATION + assert ent_reg_entry.entity_category is EntityCategory.CONFIG + + # Enable entity + entity_registry.async_update_entity( + entity_id=button_regenerate_password, disabled_by=None + ) + await hass.async_block_till_done() + + async_fire_time_changed( + hass, + dt_util.utcnow() + timedelta(seconds=RELOAD_AFTER_UPDATE_DELAY + 1), + ) + await hass.async_block_till_done() + + assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 1 + + # Validate state object + button = hass.states.get(button_regenerate_password) + assert button is not None + assert button.attributes.get(ATTR_DEVICE_CLASS) == ButtonDeviceClass.UPDATE + + aioclient_mock.clear_requests() + aioclient_mock.put( + f"https://{config_entry.data[CONF_HOST]}:1234" + f"/api/s/{config_entry.data[CONF_SITE_ID]}/rest/wlanconf/{WLAN[WLAN_ID]}", + json={"data": "password changed successfully", "meta": {"rc": "ok"}}, + headers={"content-type": CONTENT_TYPE_JSON}, + ) + + # Send WLAN regenerate password command + await hass.services.async_call( + BUTTON_DOMAIN, + "press", + {"entity_id": button_regenerate_password}, + blocking=True, + ) + assert aioclient_mock.call_count == 1 + assert next(iter(aioclient_mock.mock_calls[0][2])) == "x_passphrase" + + # Availability signalling + + # Controller disconnects + await websocket_mock.disconnect() + assert hass.states.get(button_regenerate_password).state == STATE_UNAVAILABLE + + # Controller reconnects + await websocket_mock.reconnect() + assert hass.states.get(button_regenerate_password).state != STATE_UNAVAILABLE