Add PoE power cycle button to UniFi integration

homeassistant/components/unifi/button.py

@@ -11,8 +11,14 @@
 import aiounifi
 from aiounifi.interfaces.api_handlers import ItemEvent
 from aiounifi.interfaces.devices import Devices
+from aiounifi.interfaces.ports import Ports
 from aiounifi.models.api import ApiItemT
-from aiounifi.models.device import Device, DeviceRestartRequest
+from aiounifi.models.device import (
+    Device,
+    DevicePowerCyclePortRequest,
+    DeviceRestartRequest,
+)
+from aiounifi.models.port import Port
 
 from homeassistant.components.button import (
     ButtonDeviceClass,
@@ -42,6 +48,15 @@ async def async_restart_device_control_fn(
     await api.request(DeviceRestartRequest.create(obj_id))
 
 
+@callback
+async def async_power_cycle_port_control_fn(
+    api: aiounifi.Controller, obj_id: str
+) -> None:
+    """Restart device."""
+    mac, _, index = obj_id.partition("_")
+    await api.request(DevicePowerCyclePortRequest.create(mac, int(index)))
+
+
 @dataclass
 class UnifiButtonEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
     """Validate and load entities from different UniFi handlers."""
@@ -77,6 +92,24 @@ class UnifiButtonEntityDescription(
         supported_fn=lambda controller, obj_id: True,
         unique_id_fn=lambda controller, obj_id: f"device_restart-{obj_id}",
     ),
+    UnifiButtonEntityDescription[Ports, Port](
+        key="PoE power cycle",
+        entity_category=EntityCategory.CONFIG,
+        has_entity_name=True,
+        device_class=ButtonDeviceClass.RESTART,
+        allowed_fn=lambda controller, obj_id: True,
+        api_handler_fn=lambda api: api.ports,
+        available_fn=async_device_available_fn,
+        control_fn=async_power_cycle_port_control_fn,
+        device_info_fn=async_device_device_info_fn,
+        event_is_on=None,
+        event_to_subscribe=None,
+        name_fn=lambda port: f"{port.name} Power Cycle",
+        object_fn=lambda api, obj_id: api.ports[obj_id],
+        should_poll=False,
+        supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe,
+        unique_id_fn=lambda controller, obj_id: f"power_cycle-{obj_id}",
+    ),
 )
 
 

tests/components/unifi/test_button.py

@@ -75,3 +75,89 @@ async def test_restart_device_button(
     # Controller reconnects
     await websocket_mock.reconnect()
     assert hass.states.get("button.switch_restart").state != STATE_UNAVAILABLE
+
+
+async def test_power_cycle_poe(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
+) -> None:
+    """Test restarting device button."""
+    config_entry = await setup_unifi_integration(
+        hass,
+        aioclient_mock,
+        devices_response=[
+            {
+                "board_rev": 3,
+                "device_id": "mock-id",
+                "ip": "10.0.0.1",
+                "last_seen": 1562600145,
+                "mac": "00:00:00:00:01:01",
+                "model": "US16P150",
+                "name": "switch",
+                "state": 1,
+                "type": "usw",
+                "version": "4.0.42.10433",
+                "port_table": [
+                    {
+                        "media": "GE",
+                        "name": "Port 1",
+                        "port_idx": 1,
+                        "poe_caps": 7,
+                        "poe_class": "Class 4",
+                        "poe_enable": True,
+                        "poe_mode": "auto",
+                        "poe_power": "2.56",
+                        "poe_voltage": "53.40",
+                        "portconf_id": "1a1",
+                        "port_poe": True,
+                        "up": True,
+                    },
+                ],
+            }
+        ],
+    )
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
+    assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 2
+
+    ent_reg = er.async_get(hass)
+    ent_reg_entry = ent_reg.async_get("button.switch_port_1_power_cycle")
+    assert ent_reg_entry.unique_id == "power_cycle-00:00:00:00:01:01_1"
+    assert ent_reg_entry.entity_category is EntityCategory.CONFIG
+
+    # Validate state object
+    button = hass.states.get("button.switch_port_1_power_cycle")
+    assert button is not None
+    assert button.attributes.get(ATTR_DEVICE_CLASS) == ButtonDeviceClass.RESTART
+
+    # Send restart device command
+    aioclient_mock.clear_requests()
+    aioclient_mock.post(
+        f"https://{controller.host}:1234/api/s/{controller.site}/cmd/devmgr",
+    )
+
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        "press",
+        {"entity_id": "button.switch_port_1_power_cycle"},
+        blocking=True,
+    )
+    assert aioclient_mock.call_count == 1
+    assert aioclient_mock.mock_calls[0][2] == {
+        "cmd": "power-cycle",
+        "mac": "00:00:00:00:01:01",
+        "port_idx": 1,
+    }
+
+    # Availability signalling
+
+    # Controller disconnects
+    await websocket_mock.disconnect()
+    assert (
+        hass.states.get("button.switch_port_1_power_cycle").state == STATE_UNAVAILABLE
+    )
+
+    # Controller reconnects
+    await websocket_mock.reconnect()
+    assert (
+        hass.states.get("button.switch_port_1_power_cycle").state != STATE_UNAVAILABLE
+    )