From 71bf782b22addf2179536baf2b0edca86ba128f3 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Sun, 27 Aug 2023 16:58:48 +0200 Subject: [PATCH] Improve UniFi PoE control by queueing commands together (#99114) * Working draft without timer * Clean up Improve tests * Use async_call_later --- homeassistant/components/unifi/controller.py | 35 ++++++++++++- homeassistant/components/unifi/switch.py | 52 +++++++++----------- tests/components/unifi/test_switch.py | 18 ++++++- 3 files changed, 75 insertions(+), 30 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 59cbbb5b7fd..0b0caf3add6 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -11,6 +11,7 @@ from aiohttp import CookieJar import aiounifi from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.models.configuration import Configuration +from aiounifi.models.device import DeviceSetPoePortModeRequest from aiounifi.websocket import WebsocketState from homeassistant.config_entries import ConfigEntry @@ -35,7 +36,7 @@ from homeassistant.helpers.dispatcher import ( ) from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_registry import async_entries_for_config_entry -from homeassistant.helpers.event import async_track_time_interval +from homeassistant.helpers.event import async_call_later, async_track_time_interval import homeassistant.util.dt as dt_util from .const import ( @@ -99,6 +100,9 @@ class UniFiController: self.entities: dict[str, str] = {} self.known_objects: set[tuple[str, str]] = set() + self.poe_command_queue: dict[str, dict[int, str]] = {} + self._cancel_poe_command: CALLBACK_TYPE | None = None + def load_config_entry_options(self) -> None: """Store attributes to avoid property call overhead since they are called frequently.""" options = self.config_entry.options @@ -312,6 +316,31 @@ class UniFiController: for unique_id in unique_ids_to_remove: del self._heartbeat_time[unique_id] + @callback + def async_queue_poe_port_command( + self, device_id: str, port_idx: int, poe_mode: str + ) -> None: + """Queue commands to execute them together per device.""" + if self._cancel_poe_command: + self._cancel_poe_command() + self._cancel_poe_command = None + + device_queue = self.poe_command_queue.setdefault(device_id, {}) + device_queue[port_idx] = poe_mode + + async def async_execute_command(now: datetime) -> None: + """Execute previously queued commands.""" + queue = self.poe_command_queue.copy() + self.poe_command_queue.clear() + for device_id, device_commands in queue.items(): + device = self.api.devices[device_id] + commands = [(idx, mode) for idx, mode in device_commands.items()] + await self.api.request( + DeviceSetPoePortModeRequest.create(device, targets=commands) + ) + + self._cancel_poe_command = async_call_later(self.hass, 5, async_execute_command) + async def async_update_device_registry(self) -> None: """Update device registry.""" if self.mac is None: @@ -390,6 +419,10 @@ class UniFiController: self._cancel_heartbeat_check() self._cancel_heartbeat_check = None + if self._cancel_poe_command: + self._cancel_poe_command() + self._cancel_poe_command = None + return True diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 046aa3a1abd..560e150e63c 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -22,10 +22,7 @@ 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 ( - DeviceSetOutletRelayRequest, - DeviceSetPoePortModeRequest, -) +from aiounifi.models.device import DeviceSetOutletRelayRequest from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup from aiounifi.models.event import Event, EventKey @@ -107,20 +104,22 @@ def async_port_forward_device_info_fn( async def async_block_client_control_fn( - api: aiounifi.Controller, obj_id: str, target: bool + controller: UniFiController, obj_id: str, target: bool ) -> None: """Control network access of client.""" - await api.request(ClientBlockRequest.create(obj_id, not target)) + await controller.api.request(ClientBlockRequest.create(obj_id, not target)) async def async_dpi_group_control_fn( - api: aiounifi.Controller, obj_id: str, target: bool + controller: UniFiController, obj_id: str, target: bool ) -> None: """Enable or disable DPI group.""" - dpi_group = api.dpi_groups[obj_id] + dpi_group = controller.api.dpi_groups[obj_id] await asyncio.gather( *[ - api.request(DPIRestrictionAppEnableRequest.create(app_id, target)) + controller.api.request( + DPIRestrictionAppEnableRequest.create(app_id, target) + ) for app_id in dpi_group.dpiapp_ids or [] ] ) @@ -136,46 +135,47 @@ def async_outlet_supports_switching_fn( async def async_outlet_control_fn( - api: aiounifi.Controller, obj_id: str, target: bool + controller: UniFiController, obj_id: str, target: bool ) -> None: """Control outlet relay.""" mac, _, index = obj_id.partition("_") - device = api.devices[mac] - await api.request(DeviceSetOutletRelayRequest.create(device, int(index), target)) + device = controller.api.devices[mac] + await controller.api.request( + DeviceSetOutletRelayRequest.create(device, int(index), target) + ) async def async_poe_port_control_fn( - api: aiounifi.Controller, obj_id: str, target: bool + controller: UniFiController, obj_id: str, target: bool ) -> None: """Control poe state.""" mac, _, index = obj_id.partition("_") - device = api.devices[mac] - port = api.ports[obj_id] + port = controller.api.ports[obj_id] on_state = "auto" if port.raw["poe_caps"] != 8 else "passthrough" state = on_state if target else "off" - await api.request(DeviceSetPoePortModeRequest.create(device, int(index), state)) + controller.async_queue_poe_port_command(mac, int(index), state) async def async_port_forward_control_fn( - api: aiounifi.Controller, obj_id: str, target: bool + controller: UniFiController, obj_id: str, target: bool ) -> None: """Control port forward state.""" - port_forward = api.port_forwarding[obj_id] - await api.request(PortForwardEnableRequest.create(port_forward, target)) + port_forward = controller.api.port_forwarding[obj_id] + await controller.api.request(PortForwardEnableRequest.create(port_forward, target)) async def async_wlan_control_fn( - api: aiounifi.Controller, obj_id: str, target: bool + controller: UniFiController, obj_id: str, target: bool ) -> None: """Control outlet relay.""" - await api.request(WlanEnableRequest.create(obj_id, target)) + await controller.api.request(WlanEnableRequest.create(obj_id, target)) @dataclass class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, ApiItemT]): """Validate and load entities from different UniFi handlers.""" - control_fn: Callable[[aiounifi.Controller, str, bool], Coroutine[Any, Any, None]] + control_fn: Callable[[UniFiController, str, bool], Coroutine[Any, Any, None]] is_on_fn: Callable[[UniFiController, ApiItemT], bool] @@ -352,15 +352,11 @@ class UnifiSwitchEntity(UnifiEntity[HandlerT, ApiItemT], SwitchEntity): async def async_turn_on(self, **kwargs: Any) -> None: """Turn on switch.""" - await self.entity_description.control_fn( - self.controller.api, self._obj_id, True - ) + await self.entity_description.control_fn(self.controller, self._obj_id, True) async def async_turn_off(self, **kwargs: Any) -> None: """Turn off switch.""" - await self.entity_description.control_fn( - self.controller.api, self._obj_id, False - ) + await self.entity_description.control_fn(self.controller, self._obj_id, False) @callback def async_update_state(self, event: ItemEvent, obj_id: str) -> None: diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 8e3e215e717..d376cab8add 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -1345,6 +1345,9 @@ async def test_poe_port_switches( ent_reg.async_update_entity( entity_id="switch.mock_name_port_1_poe", disabled_by=None ) + ent_reg.async_update_entity( + entity_id="switch.mock_name_port_2_poe", disabled_by=None + ) await hass.async_block_till_done() async_fire_time_changed( @@ -1378,6 +1381,8 @@ async def test_poe_port_switches( {"entity_id": "switch.mock_name_port_1_poe"}, blocking=True, ) + async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=5)) + await hass.async_block_till_done() assert aioclient_mock.call_count == 1 assert aioclient_mock.mock_calls[0][2] == { "port_overrides": [{"poe_mode": "off", "port_idx": 1, "portconf_id": "1a1"}] @@ -1390,9 +1395,20 @@ async def test_poe_port_switches( {"entity_id": "switch.mock_name_port_1_poe"}, blocking=True, ) + await hass.services.async_call( + SWITCH_DOMAIN, + "turn_off", + {"entity_id": "switch.mock_name_port_2_poe"}, + blocking=True, + ) + async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=5)) + await hass.async_block_till_done() assert aioclient_mock.call_count == 2 assert aioclient_mock.mock_calls[1][2] == { - "port_overrides": [{"poe_mode": "auto", "port_idx": 1, "portconf_id": "1a1"}] + "port_overrides": [ + {"poe_mode": "auto", "port_idx": 1, "portconf_id": "1a1"}, + {"poe_mode": "off", "port_idx": 2, "portconf_id": "1a2"}, + ] } # Availability signalling