Improve UniFi PoE control by queueing commands together (#99114)

* Working draft without timer

* Clean up
Improve tests

* Use async_call_later
This commit is contained in:
Robert Svensson 2023-08-27 16:58:48 +02:00 committed by GitHub
parent 20b8c5dd26
commit 71bf782b22
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 75 additions and 30 deletions

View File

@ -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

View File

@ -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:

View File

@ -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