diff --git a/homeassistant/components/unifi/button.py b/homeassistant/components/unifi/button.py index 7471675123a..af7ab5852ab 100644 --- a/homeassistant/components/unifi/button.py +++ b/homeassistant/components/unifi/button.py @@ -11,8 +11,14 @@ from typing import Any, Generic 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 @@ ENTITY_DESCRIPTIONS: tuple[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}", + ), ) diff --git a/tests/components/unifi/test_button.py b/tests/components/unifi/test_button.py index 30a1b3e08ff..8e6dce71160 100644 --- a/tests/components/unifi/test_button.py +++ b/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 + )