Add UniFi WLAN regenerate password button (#114422)

* Adding UniFi WLAN Change Password Button

Co-authored-by: Robert Svensson <Kane610@users.noreply.github.com>

* Adding UniFi WLAN Regenerate Password Button

Co-authored-by: Robert Svensson <Kane610@users.noreply.github.com>

---------

Co-authored-by: Robert Svensson <Kane610@users.noreply.github.com>

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]
 

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