Add support for PoE control of TP-Link Omada Gateways (#114138)

Co-authored-by: Robert Resch <robert@resch.dev>
Co-authored-by: J. Nick Koston <nick@koston.org>
This commit is contained in:
MarkGodwin 2024-03-27 15:05:14 +00:00 committed by GitHub
parent 23e9be756d
commit 834f45397d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 267 additions and 116 deletions

View File

@ -8,7 +8,6 @@ from dataclasses import dataclass
from tplink_omada_client.definitions import GatewayPortMode, LinkStatus, PoEMode from tplink_omada_client.definitions import GatewayPortMode, LinkStatus, PoEMode
from tplink_omada_client.devices import ( from tplink_omada_client.devices import (
OmadaDevice, OmadaDevice,
OmadaGateway,
OmadaGatewayPortConfig, OmadaGatewayPortConfig,
OmadaGatewayPortStatus, OmadaGatewayPortStatus,
) )
@ -95,7 +94,9 @@ GATEWAY_PORT_SENSORS: list[GatewayPortBinarySensorEntityDescription] = [
] ]
class OmadaGatewayPortBinarySensor(OmadaDeviceEntity[OmadaGateway], BinarySensorEntity): class OmadaGatewayPortBinarySensor(
OmadaDeviceEntity[OmadaGatewayCoordinator], BinarySensorEntity
):
"""Binary status of a property on an internet gateway.""" """Binary status of a property on an internet gateway."""
entity_description: GatewayPortBinarySensorEntityDescription entity_description: GatewayPortBinarySensorEntityDescription

View File

@ -1,6 +1,6 @@
"""Base entity definitions.""" """Base entity definitions."""
from typing import Generic, TypeVar from typing import Any, Generic, TypeVar
from tplink_omada_client.devices import OmadaDevice from tplink_omada_client.devices import OmadaDevice
@ -11,13 +11,13 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import DOMAIN from .const import DOMAIN
from .coordinator import OmadaCoordinator from .coordinator import OmadaCoordinator
T = TypeVar("T") T = TypeVar("T", bound="OmadaCoordinator[Any]")
class OmadaDeviceEntity(CoordinatorEntity[OmadaCoordinator[T]], Generic[T]): class OmadaDeviceEntity(CoordinatorEntity[T], Generic[T]):
"""Common base class for all entities associated with Omada SDN Devices.""" """Common base class for all entities associated with Omada SDN Devices."""
def __init__(self, coordinator: OmadaCoordinator[T], device: OmadaDevice) -> None: def __init__(self, coordinator: T, device: OmadaDevice) -> None:
"""Initialize the device.""" """Initialize the device."""
super().__init__(coordinator) super().__init__(coordinator)
self.device = device self.device = device

View File

@ -5,17 +5,19 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from dataclasses import dataclass from dataclasses import dataclass
from functools import partial from functools import partial
from typing import Any from typing import Any, Generic, TypeVar
from tplink_omada_client import OmadaSiteClient, SwitchPortOverrides from tplink_omada_client import OmadaSiteClient, SwitchPortOverrides
from tplink_omada_client.definitions import GatewayPortMode, PoEMode from tplink_omada_client.definitions import GatewayPortMode, PoEMode, PortType
from tplink_omada_client.devices import ( from tplink_omada_client.devices import (
OmadaDevice, OmadaDevice,
OmadaGateway, OmadaGateway,
OmadaGatewayPortConfig,
OmadaGatewayPortStatus, OmadaGatewayPortStatus,
OmadaSwitch, OmadaSwitch,
OmadaSwitchPortDetails, OmadaSwitchPortDetails,
) )
from tplink_omada_client.omadasiteclient import GatewayPortSettings
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
@ -29,8 +31,13 @@ from .controller import (
OmadaSiteController, OmadaSiteController,
OmadaSwitchPortCoordinator, OmadaSwitchPortCoordinator,
) )
from .coordinator import OmadaCoordinator
from .entity import OmadaDeviceEntity from .entity import OmadaDeviceEntity
TPort = TypeVar("TPort")
TDevice = TypeVar("TDevice", bound="OmadaDevice")
TCoordinator = TypeVar("TCoordinator", bound="OmadaCoordinator[Any]")
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
@ -51,37 +58,110 @@ async def async_setup_entry(
coordinator = controller.get_switch_port_coordinator(switch) coordinator = controller.get_switch_port_coordinator(switch)
await coordinator.async_request_refresh() await coordinator.async_request_refresh()
for idx, port_id in enumerate(coordinator.data): entities.extend(
if idx < switch.device_capabilities.poe_ports: OmadaDevicePortSwitchEntity[
entities.append( OmadaSwitchPortCoordinator, OmadaSwitch, OmadaSwitchPortDetails
OmadaNetworkSwitchPortPoEControl(coordinator, switch, port_id) ](
) coordinator,
switch,
port.port_id,
desc,
port_name=_get_switch_port_base_name(port),
)
for port in coordinator.data.values()
for desc in SWITCH_PORT_DETAILS_SWITCHES
if desc.exists_func(switch, port)
)
gateway_coordinator = await controller.get_gateway_coordinator() gateway_coordinator = await controller.get_gateway_coordinator()
if gateway_coordinator: if gateway_coordinator:
for gateway in gateway_coordinator.data.values(): for gateway in gateway_coordinator.data.values():
entities.extend( entities.extend(
OmadaGatewayPortSwitchEntity( OmadaDevicePortSwitchEntity[
gateway_coordinator, gateway, p.port_number, desc OmadaGatewayCoordinator, OmadaGateway, OmadaGatewayPortStatus
) ](gateway_coordinator, gateway, p.port_number, desc)
for p in gateway.port_status for p in gateway.port_status
for desc in GATEWAY_PORT_SWITCHES for desc in GATEWAY_PORT_STATUS_SWITCHES
if desc.exists_func(p) if desc.exists_func(gateway, p)
)
entities.extend(
OmadaDevicePortSwitchEntity[
OmadaGatewayCoordinator, OmadaGateway, OmadaGatewayPortConfig
](gateway_coordinator, gateway, p.port_number, desc)
for p in gateway.port_configs
for desc in GATEWAY_PORT_CONFIG_SWITCHES
if desc.exists_func(gateway, p)
) )
async_add_entities(entities) async_add_entities(entities)
@dataclass(frozen=True, kw_only=True) def _get_switch_port_base_name(port: OmadaSwitchPortDetails) -> str:
class GatewayPortSwitchEntityDescription(SwitchEntityDescription): """Get display name for a switch port."""
"""Entity description for a toggle switch derived from a gateway port."""
exists_func: Callable[[OmadaGatewayPortStatus], bool] = lambda _: True if port.name == f"Port{port.port}":
set_func: Callable[ return str(port.port)
[OmadaSiteClient, OmadaDevice, OmadaGatewayPortStatus, bool], return f"{port.port} ({port.name})"
Awaitable[OmadaGatewayPortStatus],
@dataclass(frozen=True, kw_only=True)
class OmadaDevicePortSwitchEntityDescription(
SwitchEntityDescription, Generic[TCoordinator, TDevice, TPort]
):
"""Entity description for a toggle switch derived from a network port on an Omada device."""
exists_func: Callable[[TDevice, TPort], bool] = lambda _, p: True
coordinator_update_func: Callable[
[TCoordinator, TDevice, int | str], TPort | None
] = lambda *_: None
set_func: Callable[[OmadaSiteClient, TDevice, TPort, bool], Awaitable[TPort]]
update_func: Callable[[TPort], bool]
refresh_after_set: bool = False
@dataclass(frozen=True, kw_only=True)
class OmadaSwitchPortSwitchEntityDescription(
OmadaDevicePortSwitchEntityDescription[
OmadaSwitchPortCoordinator, OmadaSwitch, OmadaSwitchPortDetails
] ]
update_func: Callable[[OmadaGatewayPortStatus], bool] ):
"""Entity description for a toggle switch for a feature of a Port on an Omada Switch."""
coordinator_update_func: Callable[
[OmadaSwitchPortCoordinator, OmadaSwitch, int | str],
OmadaSwitchPortDetails | None,
] = lambda coord, _, port_id: coord.data.get(str(port_id))
@dataclass(frozen=True, kw_only=True)
class OmadaGatewayPortConfigSwitchEntityDescription(
OmadaDevicePortSwitchEntityDescription[
OmadaGatewayCoordinator, OmadaGateway, OmadaGatewayPortConfig
]
):
"""Entity description for a toggle switch for a configuration of a Port on an Omada Gateway."""
coordinator_update_func: Callable[
[OmadaGatewayCoordinator, OmadaGateway, int | str],
OmadaGatewayPortConfig | None,
] = lambda coord, device, port_id: next(
p for p in coord.data[device.mac].port_configs if p.port_number == port_id
)
@dataclass(frozen=True, kw_only=True)
class OmadaGatewayPortStatusSwitchEntityDescription(
OmadaDevicePortSwitchEntityDescription[
OmadaGatewayCoordinator, OmadaGateway, OmadaGatewayPortStatus
]
):
"""Entity description for a toggle switch for a status of a Port on an Omada Gateway."""
coordinator_update_func: Callable[
[OmadaGatewayCoordinator, OmadaGateway, int | str], OmadaGatewayPortStatus
] = lambda coord, device, port_id: next(
p for p in coord.data[device.mac].port_status if p.port_number == port_id
)
def _wan_connect_disconnect( def _wan_connect_disconnect(
@ -96,109 +176,82 @@ def _wan_connect_disconnect(
) )
GATEWAY_PORT_SWITCHES: list[GatewayPortSwitchEntityDescription] = [ SWITCH_PORT_DETAILS_SWITCHES: list[OmadaSwitchPortSwitchEntityDescription] = [
GatewayPortSwitchEntityDescription( OmadaSwitchPortSwitchEntityDescription(
key="poe",
translation_key="poe_control",
exists_func=lambda d, p: d.device_capabilities.supports_poe
and p.type != PortType.SFP,
set_func=lambda client, device, port, enable: client.update_switch_port(
device, port, overrides=SwitchPortOverrides(enable_poe=enable)
),
update_func=lambda p: p.poe_mode != PoEMode.DISABLED,
entity_category=EntityCategory.CONFIG,
)
]
GATEWAY_PORT_STATUS_SWITCHES: list[OmadaGatewayPortStatusSwitchEntityDescription] = [
OmadaGatewayPortStatusSwitchEntityDescription(
key="wan_connect_ipv4", key="wan_connect_ipv4",
translation_key="wan_connect_ipv4", translation_key="wan_connect_ipv4",
exists_func=lambda p: p.mode == GatewayPortMode.WAN, exists_func=lambda _, p: p.mode == GatewayPortMode.WAN,
set_func=partial(_wan_connect_disconnect, ipv6=False), set_func=partial(_wan_connect_disconnect, ipv6=False),
update_func=lambda p: p.wan_connected, update_func=lambda p: p.wan_connected,
refresh_after_set=True,
), ),
GatewayPortSwitchEntityDescription( OmadaGatewayPortStatusSwitchEntityDescription(
key="wan_connect_ipv6", key="wan_connect_ipv6",
translation_key="wan_connect_ipv6", translation_key="wan_connect_ipv6",
exists_func=lambda p: p.mode == GatewayPortMode.WAN and p.wan_ipv6_enabled, exists_func=lambda _, p: p.mode == GatewayPortMode.WAN and p.wan_ipv6_enabled,
set_func=partial(_wan_connect_disconnect, ipv6=True), set_func=partial(_wan_connect_disconnect, ipv6=True),
update_func=lambda p: p.ipv6_wan_connected, update_func=lambda p: p.ipv6_wan_connected,
refresh_after_set=True,
),
]
GATEWAY_PORT_CONFIG_SWITCHES: list[OmadaGatewayPortConfigSwitchEntityDescription] = [
OmadaGatewayPortConfigSwitchEntityDescription(
key="poe",
translation_key="poe_control",
exists_func=lambda _, port: port.poe_mode != PoEMode.NONE,
set_func=lambda client, device, port, enable: client.set_gateway_port_settings(
port.port_number, GatewayPortSettings(enable_poe=enable), device
),
update_func=lambda p: p.poe_mode != PoEMode.DISABLED,
), ),
] ]
def get_port_base_name(port: OmadaSwitchPortDetails) -> str: class OmadaDevicePortSwitchEntity(
"""Get display name for a switch port.""" OmadaDeviceEntity[TCoordinator],
SwitchEntity,
if port.name == f"Port{port.port}": Generic[TCoordinator, TDevice, TPort],
return f"{port.port}"
return f"{port.port} ({port.name})"
class OmadaNetworkSwitchPortPoEControl(
OmadaDeviceEntity[OmadaSwitchPortDetails], SwitchEntity
): ):
"""Representation of a PoE control toggle on a single network port on a switch.""" """Generic toggle switch entity for a Netork Port of an Omada Device."""
_attr_has_entity_name = True _attr_has_entity_name = True
_attr_translation_key = "poe_control" _port_details: TPort | None = None
_attr_entity_category = EntityCategory.CONFIG entity_description: OmadaDevicePortSwitchEntityDescription[
TCoordinator, TDevice, TPort
]
def __init__( def __init__(
self, self,
coordinator: OmadaSwitchPortCoordinator, coordinator: TCoordinator,
device: OmadaSwitch, device: TDevice,
port_id: str, port_id: int | str,
) -> None: entity_description: OmadaDevicePortSwitchEntityDescription[
"""Initialize the PoE switch.""" TCoordinator, TDevice, TPort
super().__init__(coordinator, device) ],
self.port_id = port_id port_name: str | None = None,
self.port_details = coordinator.data[port_id]
self.omada_client = coordinator.omada_client
self._attr_unique_id = f"{device.mac}_{port_id}_poe"
self._attr_translation_placeholders = {
"port_name": get_port_base_name(self.port_details)
}
async def async_added_to_hass(self) -> None:
"""When entity is added to hass."""
await super().async_added_to_hass()
self._refresh_state()
async def _async_turn_on_off_poe(self, enable: bool) -> None:
self.port_details = await self.omada_client.update_switch_port(
self.device,
self.port_details,
overrides=SwitchPortOverrides(enable_poe=enable),
)
self._refresh_state()
async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn the entity on."""
await self._async_turn_on_off_poe(True)
async def async_turn_off(self, **kwargs: Any) -> None:
"""Turn the entity off."""
await self._async_turn_on_off_poe(False)
def _refresh_state(self) -> None:
self._attr_is_on = self.port_details.poe_mode != PoEMode.DISABLED
self.async_write_ha_state()
@callback
def _handle_coordinator_update(self) -> None:
"""Handle updated data from the coordinator."""
self.port_details = self.coordinator.data[self.port_id]
self._refresh_state()
class OmadaGatewayPortSwitchEntity(OmadaDeviceEntity[OmadaGateway], SwitchEntity):
"""Generic toggle switch on a Gateway entity."""
_attr_has_entity_name = True
_port_details: OmadaGatewayPortStatus | None = None
entity_description: GatewayPortSwitchEntityDescription
def __init__(
self,
coordinator: OmadaGatewayCoordinator,
device: OmadaGateway,
port_number: int,
entity_description: GatewayPortSwitchEntityDescription,
) -> None: ) -> None:
"""Initialize the toggle switch.""" """Initialize the toggle switch."""
super().__init__(coordinator, device) super().__init__(coordinator, device)
self.entity_description = entity_description self.entity_description = entity_description
self._port_number = port_number self._device = device
self._attr_unique_id = f"{device.mac}_{port_number}_{entity_description.key}" self._port_id = port_id
self._attr_translation_placeholders = {"port_name": f"{port_number}"} self._attr_unique_id = f"{device.mac}_{port_id}_{entity_description.key}"
self._attr_translation_placeholders = {"port_name": port_name or str(port_id)}
async def async_added_to_hass(self) -> None: async def async_added_to_hass(self) -> None:
"""When entity is added to hass.""" """When entity is added to hass."""
@ -208,11 +261,16 @@ class OmadaGatewayPortSwitchEntity(OmadaDeviceEntity[OmadaGateway], SwitchEntity
async def _async_turn_on_off(self, enable: bool) -> None: async def _async_turn_on_off(self, enable: bool) -> None:
if self._port_details: if self._port_details:
self._port_details = await self.entity_description.set_func( self._port_details = await self.entity_description.set_func(
self.coordinator.omada_client, self.device, self._port_details, enable self.coordinator.omada_client, self._device, self._port_details, enable
) )
self._attr_is_on = enable
# Refresh to make sure the requested changes stuck if self.entity_description.refresh_after_set:
await self.coordinator.async_request_refresh() # Refresh to make sure the requested changes stuck
self._attr_is_on = enable
await self.coordinator.async_request_refresh()
elif self._port_details:
self._attr_is_on = self.entity_description.update_func(self._port_details)
self.async_write_ha_state()
async def async_turn_on(self, **kwargs: Any) -> None: async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn the entity on.""" """Turn the entity on."""
@ -228,14 +286,12 @@ class OmadaGatewayPortSwitchEntity(OmadaDeviceEntity[OmadaGateway], SwitchEntity
return bool( return bool(
super().available super().available
and self._port_details and self._port_details
and self.entity_description.exists_func(self._port_details) and self.entity_description.exists_func(self._device, self._port_details)
) )
def _do_update(self) -> None: def _do_update(self) -> None:
gateway = self.coordinator.data[self.device.mac] port = self.entity_description.coordinator_update_func(
self.coordinator, self._device, self._port_id
port = next(
p for p in gateway.port_status if p.port_number == self._port_number
) )
if port: if port:
self._port_details = port self._port_details = port

View File

@ -88,7 +88,7 @@ async def async_setup_entry(
class OmadaDeviceUpdate( class OmadaDeviceUpdate(
OmadaDeviceEntity[FirmwareUpdateStatus], OmadaDeviceEntity[OmadaFirmwareUpdateCoodinator],
UpdateEntity, UpdateEntity,
): ):
"""Firmware update status for Omada SDN devices.""" """Firmware update status for Omada SDN devices."""

View File

@ -51,6 +51,19 @@
'state': 'on', 'state': 'on',
}) })
# --- # ---
# name: test_gateway_port_poe_switch
StateSnapshot({
'attributes': ReadOnlyDict({
'friendly_name': 'Test Router Port 5 PoE',
}),
'context': <ANY>,
'entity_id': 'switch.test_router_port_5_poe',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'on',
})
# ---
# name: test_poe_switches # name: test_poe_switches
StateSnapshot({ StateSnapshot({
'attributes': ReadOnlyDict({ 'attributes': ReadOnlyDict({

View File

@ -9,6 +9,7 @@ from tplink_omada_client import SwitchPortOverrides
from tplink_omada_client.definitions import PoEMode from tplink_omada_client.definitions import PoEMode
from tplink_omada_client.devices import ( from tplink_omada_client.devices import (
OmadaGateway, OmadaGateway,
OmadaGatewayPortConfig,
OmadaGatewayPortStatus, OmadaGatewayPortStatus,
OmadaSwitch, OmadaSwitch,
OmadaSwitchPortDetails, OmadaSwitchPortDetails,
@ -55,6 +56,17 @@ async def test_poe_switches(
) )
async def test_sfp_port_has_no_poe_switch(
hass: HomeAssistant,
init_integration: MockConfigEntry,
) -> None:
"""Test PoE switch SFP ports have no PoE controls."""
entity = hass.states.get("switch.test_poe_switch_port_9_poe")
assert entity is None
entity = hass.states.get("switch.test_poe_switch_port_8_poe")
assert entity is not None
async def test_gateway_connect_ipv4_switch( async def test_gateway_connect_ipv4_switch(
hass: HomeAssistant, hass: HomeAssistant,
mock_omada_site_client: MagicMock, mock_omada_site_client: MagicMock,
@ -107,6 +119,70 @@ async def test_gateway_connect_ipv4_switch(
assert entity.state == "on" assert entity.state == "on"
async def test_gateway_port_poe_switch(
hass: HomeAssistant,
mock_omada_site_client: MagicMock,
init_integration: MockConfigEntry,
snapshot: SnapshotAssertion,
) -> None:
"""Test gateway connected switches."""
gateway_mac = "AA-BB-CC-DD-EE-FF"
entity_id = "switch.test_router_port_5_poe"
entity = hass.states.get(entity_id)
assert entity == snapshot
test_gateway = await mock_omada_site_client.get_gateway(gateway_mac)
port_config = test_gateway.port_configs[4]
assert port_config.port_number == 5
mock_omada_site_client.set_gateway_port_settings.return_value = (
OmadaGatewayPortConfig(port_config.raw_data, poe_enabled=False)
)
await call_service(hass, "turn_off", entity_id)
_assert_gateway_poe_set(mock_omada_site_client, test_gateway, False)
async_fire_time_changed(hass, utcnow() + UPDATE_INTERVAL)
await hass.async_block_till_done()
entity = hass.states.get(entity_id)
assert entity.state == "off"
mock_omada_site_client.set_gateway_port_settings.reset_mock()
mock_omada_site_client.set_gateway_port_settings.return_value = port_config
await call_service(hass, "turn_on", entity_id)
_assert_gateway_poe_set(mock_omada_site_client, test_gateway, True)
async_fire_time_changed(hass, utcnow() + UPDATE_INTERVAL)
await hass.async_block_till_done()
entity = hass.states.get(entity_id)
assert entity.state == "on"
async def test_gaateway_wan_port_has_no_poe_switch(
hass: HomeAssistant,
init_integration: MockConfigEntry,
) -> None:
"""Test PoE switch SFP ports have no PoE controls."""
entity = hass.states.get("switch.test_router_port_1_poe")
assert entity is None
entity = hass.states.get("switch.test_router_port_9_poe")
assert entity is not None
def _assert_gateway_poe_set(mock_omada_site_client, test_gateway, poe_enabled: bool):
(
called_port,
called_settings,
called_gateway,
) = mock_omada_site_client.set_gateway_port_settings.call_args.args
mock_omada_site_client.set_gateway_port_settings.assert_called_once()
assert called_port == 5
assert called_settings.enable_poe is poe_enabled
assert called_gateway == test_gateway
async def test_gateway_api_fail_disables_switch_entities( async def test_gateway_api_fail_disables_switch_entities(
hass: HomeAssistant, hass: HomeAssistant,
mock_omada_site_client: MagicMock, mock_omada_site_client: MagicMock,
@ -188,18 +264,22 @@ async def _test_poe_switch(
mock_omada_site_client.update_switch_port.return_value = await _update_port_details( mock_omada_site_client.update_switch_port.return_value = await _update_port_details(
mock_omada_site_client, port_num, False mock_omada_site_client, port_num, False
) )
await call_service(hass, "turn_off", entity_id) await call_service(hass, "turn_off", entity_id)
mock_omada_site_client.update_switch_port.assert_called_once() mock_omada_site_client.update_switch_port.assert_called_once()
( (
device, device,
switch_port_details, switch_port_details,
) = mock_omada_site_client.update_switch_port.call_args.args ) = mock_omada_site_client.update_switch_port.call_args.args
assert_update_switch_port( assert_update_switch_port(
device, device,
switch_port_details, switch_port_details,
False, False,
**mock_omada_site_client.update_switch_port.call_args.kwargs, **mock_omada_site_client.update_switch_port.call_args.kwargs,
) )
async_fire_time_changed(hass, utcnow() + UPDATE_INTERVAL)
await hass.async_block_till_done()
entity = hass.states.get(entity_id) entity = hass.states.get(entity_id)
assert entity.state == "off" assert entity.state == "off"
@ -216,6 +296,7 @@ async def _test_poe_switch(
True, True,
**mock_omada_site_client.update_switch_port.call_args.kwargs, **mock_omada_site_client.update_switch_port.call_args.kwargs,
) )
async_fire_time_changed(hass, utcnow() + UPDATE_INTERVAL)
await hass.async_block_till_done() await hass.async_block_till_done()
entity = hass.states.get(entity_id) entity = hass.states.get(entity_id)
assert entity.state == "on" assert entity.state == "on"