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 import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.models.configuration import Configuration from aiounifi.models.configuration import Configuration
from aiounifi.models.device import DeviceSetPoePortModeRequest
from aiounifi.websocket import WebsocketState from aiounifi.websocket import WebsocketState
from homeassistant.config_entries import ConfigEntry 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_platform import AddEntitiesCallback
from homeassistant.helpers.entity_registry import async_entries_for_config_entry 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 import homeassistant.util.dt as dt_util
from .const import ( from .const import (
@ -99,6 +100,9 @@ class UniFiController:
self.entities: dict[str, str] = {} self.entities: dict[str, str] = {}
self.known_objects: set[tuple[str, str]] = set() 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: def load_config_entry_options(self) -> None:
"""Store attributes to avoid property call overhead since they are called frequently.""" """Store attributes to avoid property call overhead since they are called frequently."""
options = self.config_entry.options options = self.config_entry.options
@ -312,6 +316,31 @@ class UniFiController:
for unique_id in unique_ids_to_remove: for unique_id in unique_ids_to_remove:
del self._heartbeat_time[unique_id] 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: async def async_update_device_registry(self) -> None:
"""Update device registry.""" """Update device registry."""
if self.mac is None: if self.mac is None:
@ -390,6 +419,10 @@ class UniFiController:
self._cancel_heartbeat_check() self._cancel_heartbeat_check()
self._cancel_heartbeat_check = None self._cancel_heartbeat_check = None
if self._cancel_poe_command:
self._cancel_poe_command()
self._cancel_poe_command = None
return True return True

View File

@ -22,10 +22,7 @@ from aiounifi.interfaces.ports import Ports
from aiounifi.interfaces.wlans import Wlans from aiounifi.interfaces.wlans import Wlans
from aiounifi.models.api import ApiItemT from aiounifi.models.api import ApiItemT
from aiounifi.models.client import Client, ClientBlockRequest from aiounifi.models.client import Client, ClientBlockRequest
from aiounifi.models.device import ( from aiounifi.models.device import DeviceSetOutletRelayRequest
DeviceSetOutletRelayRequest,
DeviceSetPoePortModeRequest,
)
from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest
from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup
from aiounifi.models.event import Event, EventKey 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( async def async_block_client_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool controller: UniFiController, obj_id: str, target: bool
) -> None: ) -> None:
"""Control network access of client.""" """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( async def async_dpi_group_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool controller: UniFiController, obj_id: str, target: bool
) -> None: ) -> None:
"""Enable or disable DPI group.""" """Enable or disable DPI group."""
dpi_group = api.dpi_groups[obj_id] dpi_group = controller.api.dpi_groups[obj_id]
await asyncio.gather( 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 [] 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( async def async_outlet_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool controller: UniFiController, obj_id: str, target: bool
) -> None: ) -> None:
"""Control outlet relay.""" """Control outlet relay."""
mac, _, index = obj_id.partition("_") mac, _, index = obj_id.partition("_")
device = api.devices[mac] device = controller.api.devices[mac]
await api.request(DeviceSetOutletRelayRequest.create(device, int(index), target)) await controller.api.request(
DeviceSetOutletRelayRequest.create(device, int(index), target)
)
async def async_poe_port_control_fn( async def async_poe_port_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool controller: UniFiController, obj_id: str, target: bool
) -> None: ) -> None:
"""Control poe state.""" """Control poe state."""
mac, _, index = obj_id.partition("_") mac, _, index = obj_id.partition("_")
device = api.devices[mac] port = controller.api.ports[obj_id]
port = api.ports[obj_id]
on_state = "auto" if port.raw["poe_caps"] != 8 else "passthrough" on_state = "auto" if port.raw["poe_caps"] != 8 else "passthrough"
state = on_state if target else "off" 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( async def async_port_forward_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool controller: UniFiController, obj_id: str, target: bool
) -> None: ) -> None:
"""Control port forward state.""" """Control port forward state."""
port_forward = api.port_forwarding[obj_id] port_forward = controller.api.port_forwarding[obj_id]
await api.request(PortForwardEnableRequest.create(port_forward, target)) await controller.api.request(PortForwardEnableRequest.create(port_forward, target))
async def async_wlan_control_fn( async def async_wlan_control_fn(
api: aiounifi.Controller, obj_id: str, target: bool controller: UniFiController, obj_id: str, target: bool
) -> None: ) -> None:
"""Control outlet relay.""" """Control outlet relay."""
await api.request(WlanEnableRequest.create(obj_id, target)) await controller.api.request(WlanEnableRequest.create(obj_id, target))
@dataclass @dataclass
class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, ApiItemT]): class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
"""Validate and load entities from different UniFi handlers.""" """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] 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: async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn on switch.""" """Turn on switch."""
await self.entity_description.control_fn( await self.entity_description.control_fn(self.controller, self._obj_id, True)
self.controller.api, self._obj_id, True
)
async def async_turn_off(self, **kwargs: Any) -> None: async def async_turn_off(self, **kwargs: Any) -> None:
"""Turn off switch.""" """Turn off switch."""
await self.entity_description.control_fn( await self.entity_description.control_fn(self.controller, self._obj_id, False)
self.controller.api, self._obj_id, False
)
@callback @callback
def async_update_state(self, event: ItemEvent, obj_id: str) -> None: 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( ent_reg.async_update_entity(
entity_id="switch.mock_name_port_1_poe", disabled_by=None 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() await hass.async_block_till_done()
async_fire_time_changed( async_fire_time_changed(
@ -1378,6 +1381,8 @@ async def test_poe_port_switches(
{"entity_id": "switch.mock_name_port_1_poe"}, {"entity_id": "switch.mock_name_port_1_poe"},
blocking=True, 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.call_count == 1
assert aioclient_mock.mock_calls[0][2] == { assert aioclient_mock.mock_calls[0][2] == {
"port_overrides": [{"poe_mode": "off", "port_idx": 1, "portconf_id": "1a1"}] "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"}, {"entity_id": "switch.mock_name_port_1_poe"},
blocking=True, 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.call_count == 2
assert aioclient_mock.mock_calls[1][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 # Availability signalling