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>

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

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

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

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."""

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': <ANY>,
+    'entity_id': 'switch.test_router_port_5_poe',
+    'last_changed': <ANY>,
+    'last_reported': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on',
+  })
+# ---
 # name: test_poe_switches
   StateSnapshot({
     'attributes': ReadOnlyDict({

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"