From 834f45397d6a6e3c8d350f635f41bd2843730b7e Mon Sep 17 00:00:00 2001 From: MarkGodwin <10632972+MarkGodwin@users.noreply.github.com> Date: Wed, 27 Mar 2024 15:05:14 +0000 Subject: [PATCH] Add support for PoE control of TP-Link Omada Gateways (#114138) Co-authored-by: Robert Resch Co-authored-by: J. Nick Koston --- .../components/tplink_omada/binary_sensor.py | 5 +- .../components/tplink_omada/entity.py | 8 +- .../components/tplink_omada/switch.py | 274 +++++++++++------- .../components/tplink_omada/update.py | 2 +- .../tplink_omada/snapshots/test_switch.ambr | 13 + tests/components/tplink_omada/test_switch.py | 81 ++++++ 6 files changed, 267 insertions(+), 116 deletions(-) diff --git a/homeassistant/components/tplink_omada/binary_sensor.py b/homeassistant/components/tplink_omada/binary_sensor.py index 7bee6159dd7..c0304c4d1b2 100644 --- a/homeassistant/components/tplink_omada/binary_sensor.py +++ b/homeassistant/components/tplink_omada/binary_sensor.py @@ -8,7 +8,6 @@ from dataclasses import dataclass from tplink_omada_client.definitions import GatewayPortMode, LinkStatus, PoEMode from tplink_omada_client.devices import ( OmadaDevice, - OmadaGateway, OmadaGatewayPortConfig, 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.""" entity_description: GatewayPortBinarySensorEntityDescription diff --git a/homeassistant/components/tplink_omada/entity.py b/homeassistant/components/tplink_omada/entity.py index 4ae9dc733d8..a0bb562c652 100644 --- a/homeassistant/components/tplink_omada/entity.py +++ b/homeassistant/components/tplink_omada/entity.py @@ -1,6 +1,6 @@ """Base entity definitions.""" -from typing import Generic, TypeVar +from typing import Any, Generic, TypeVar from tplink_omada_client.devices import OmadaDevice @@ -11,13 +11,13 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN 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.""" - def __init__(self, coordinator: OmadaCoordinator[T], device: OmadaDevice) -> None: + def __init__(self, coordinator: T, device: OmadaDevice) -> None: """Initialize the device.""" super().__init__(coordinator) self.device = device diff --git a/homeassistant/components/tplink_omada/switch.py b/homeassistant/components/tplink_omada/switch.py index 30974e829e2..b8abb4cb773 100644 --- a/homeassistant/components/tplink_omada/switch.py +++ b/homeassistant/components/tplink_omada/switch.py @@ -5,17 +5,19 @@ from __future__ import annotations from collections.abc import Awaitable, Callable from dataclasses import dataclass 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.definitions import GatewayPortMode, PoEMode +from tplink_omada_client.definitions import GatewayPortMode, PoEMode, PortType from tplink_omada_client.devices import ( OmadaDevice, OmadaGateway, + OmadaGatewayPortConfig, OmadaGatewayPortStatus, OmadaSwitch, OmadaSwitchPortDetails, ) +from tplink_omada_client.omadasiteclient import GatewayPortSettings from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription from homeassistant.config_entries import ConfigEntry @@ -29,8 +31,13 @@ from .controller import ( OmadaSiteController, OmadaSwitchPortCoordinator, ) +from .coordinator import OmadaCoordinator from .entity import OmadaDeviceEntity +TPort = TypeVar("TPort") +TDevice = TypeVar("TDevice", bound="OmadaDevice") +TCoordinator = TypeVar("TCoordinator", bound="OmadaCoordinator[Any]") + async def async_setup_entry( hass: HomeAssistant, @@ -51,37 +58,110 @@ async def async_setup_entry( coordinator = controller.get_switch_port_coordinator(switch) await coordinator.async_request_refresh() - for idx, port_id in enumerate(coordinator.data): - if idx < switch.device_capabilities.poe_ports: - entities.append( - OmadaNetworkSwitchPortPoEControl(coordinator, switch, port_id) - ) + entities.extend( + OmadaDevicePortSwitchEntity[ + OmadaSwitchPortCoordinator, OmadaSwitch, OmadaSwitchPortDetails + ]( + 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() if gateway_coordinator: for gateway in gateway_coordinator.data.values(): entities.extend( - OmadaGatewayPortSwitchEntity( - gateway_coordinator, gateway, p.port_number, desc - ) + OmadaDevicePortSwitchEntity[ + OmadaGatewayCoordinator, OmadaGateway, OmadaGatewayPortStatus + ](gateway_coordinator, gateway, p.port_number, desc) for p in gateway.port_status - for desc in GATEWAY_PORT_SWITCHES - if desc.exists_func(p) + for desc in GATEWAY_PORT_STATUS_SWITCHES + 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) -@dataclass(frozen=True, kw_only=True) -class GatewayPortSwitchEntityDescription(SwitchEntityDescription): - """Entity description for a toggle switch derived from a gateway port.""" +def _get_switch_port_base_name(port: OmadaSwitchPortDetails) -> str: + """Get display name for a switch port.""" - exists_func: Callable[[OmadaGatewayPortStatus], bool] = lambda _: True - set_func: Callable[ - [OmadaSiteClient, OmadaDevice, OmadaGatewayPortStatus, bool], - Awaitable[OmadaGatewayPortStatus], + if port.name == f"Port{port.port}": + return str(port.port) + return f"{port.port} ({port.name})" + + +@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( @@ -96,109 +176,82 @@ def _wan_connect_disconnect( ) -GATEWAY_PORT_SWITCHES: list[GatewayPortSwitchEntityDescription] = [ - GatewayPortSwitchEntityDescription( +SWITCH_PORT_DETAILS_SWITCHES: list[OmadaSwitchPortSwitchEntityDescription] = [ + 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", 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), update_func=lambda p: p.wan_connected, + refresh_after_set=True, ), - GatewayPortSwitchEntityDescription( + OmadaGatewayPortStatusSwitchEntityDescription( 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), 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: - """Get display name for a switch port.""" - - if port.name == f"Port{port.port}": - return f"{port.port}" - return f"{port.port} ({port.name})" - - -class OmadaNetworkSwitchPortPoEControl( - OmadaDeviceEntity[OmadaSwitchPortDetails], SwitchEntity +class OmadaDevicePortSwitchEntity( + OmadaDeviceEntity[TCoordinator], + SwitchEntity, + Generic[TCoordinator, TDevice, TPort], ): - """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_translation_key = "poe_control" - _attr_entity_category = EntityCategory.CONFIG + _port_details: TPort | None = None + entity_description: OmadaDevicePortSwitchEntityDescription[ + TCoordinator, TDevice, TPort + ] def __init__( self, - coordinator: OmadaSwitchPortCoordinator, - device: OmadaSwitch, - port_id: str, - ) -> None: - """Initialize the PoE switch.""" - super().__init__(coordinator, device) - self.port_id = port_id - 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, + coordinator: TCoordinator, + device: TDevice, + port_id: int | str, + entity_description: OmadaDevicePortSwitchEntityDescription[ + TCoordinator, TDevice, TPort + ], + port_name: str | None = None, ) -> None: """Initialize the toggle switch.""" super().__init__(coordinator, device) self.entity_description = entity_description - self._port_number = port_number - self._attr_unique_id = f"{device.mac}_{port_number}_{entity_description.key}" - self._attr_translation_placeholders = {"port_name": f"{port_number}"} + self._device = device + self._port_id = port_id + 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: """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: if self._port_details: 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 - await self.coordinator.async_request_refresh() + + if self.entity_description.refresh_after_set: + # 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: """Turn the entity on.""" @@ -228,14 +286,12 @@ class OmadaGatewayPortSwitchEntity(OmadaDeviceEntity[OmadaGateway], SwitchEntity return bool( super().available 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: - gateway = self.coordinator.data[self.device.mac] - - port = next( - p for p in gateway.port_status if p.port_number == self._port_number + port = self.entity_description.coordinator_update_func( + self.coordinator, self._device, self._port_id ) if port: self._port_details = port diff --git a/homeassistant/components/tplink_omada/update.py b/homeassistant/components/tplink_omada/update.py index 8a0d32bda18..5e87d11474b 100644 --- a/homeassistant/components/tplink_omada/update.py +++ b/homeassistant/components/tplink_omada/update.py @@ -88,7 +88,7 @@ async def async_setup_entry( class OmadaDeviceUpdate( - OmadaDeviceEntity[FirmwareUpdateStatus], + OmadaDeviceEntity[OmadaFirmwareUpdateCoodinator], UpdateEntity, ): """Firmware update status for Omada SDN devices.""" diff --git a/tests/components/tplink_omada/snapshots/test_switch.ambr b/tests/components/tplink_omada/snapshots/test_switch.ambr index 8a08cbc292d..282d2a4a6a5 100644 --- a/tests/components/tplink_omada/snapshots/test_switch.ambr +++ b/tests/components/tplink_omada/snapshots/test_switch.ambr @@ -51,6 +51,19 @@ 'state': 'on', }) # --- +# name: test_gateway_port_poe_switch + StateSnapshot({ + 'attributes': ReadOnlyDict({ + 'friendly_name': 'Test Router Port 5 PoE', + }), + 'context': , + 'entity_id': 'switch.test_router_port_5_poe', + 'last_changed': , + 'last_reported': , + 'last_updated': , + 'state': 'on', + }) +# --- # name: test_poe_switches StateSnapshot({ 'attributes': ReadOnlyDict({ diff --git a/tests/components/tplink_omada/test_switch.py b/tests/components/tplink_omada/test_switch.py index bcc727da06e..78b22a4e829 100644 --- a/tests/components/tplink_omada/test_switch.py +++ b/tests/components/tplink_omada/test_switch.py @@ -9,6 +9,7 @@ from tplink_omada_client import SwitchPortOverrides from tplink_omada_client.definitions import PoEMode from tplink_omada_client.devices import ( OmadaGateway, + OmadaGatewayPortConfig, OmadaGatewayPortStatus, OmadaSwitch, 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( hass: HomeAssistant, mock_omada_site_client: MagicMock, @@ -107,6 +119,70 @@ async def test_gateway_connect_ipv4_switch( 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( hass: HomeAssistant, 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, port_num, False ) + await call_service(hass, "turn_off", entity_id) mock_omada_site_client.update_switch_port.assert_called_once() ( device, switch_port_details, ) = mock_omada_site_client.update_switch_port.call_args.args + assert_update_switch_port( device, switch_port_details, False, **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) assert entity.state == "off" @@ -216,6 +296,7 @@ async def _test_poe_switch( True, **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) assert entity.state == "on"