Add Update entities to TP-Link Omada integration (#89562)

* Bump tplink-omada

* Add omada firmware updates

* Excluded from code coverage

* Fixed entity name

.coveragerc

@@ -1290,9 +1290,11 @@ omit =
     homeassistant/components/touchline/climate.py
     homeassistant/components/tplink_lte/*
     homeassistant/components/tplink_omada/__init__.py
+    homeassistant/components/tplink_omada/controller.py
     homeassistant/components/tplink_omada/coordinator.py
     homeassistant/components/tplink_omada/entity.py
     homeassistant/components/tplink_omada/switch.py
+    homeassistant/components/tplink_omada/update.py
     homeassistant/components/traccar/device_tracker.py
     homeassistant/components/tractive/__init__.py
     homeassistant/components/tractive/binary_sensor.py

homeassistant/components/tplink_omada/__init__.py

@@ -16,8 +16,9 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 
 from .config_flow import CONF_SITE, create_omada_client
 from .const import DOMAIN
+from .controller import OmadaSiteController
 
-PLATFORMS: list[Platform] = [Platform.SWITCH]
+PLATFORMS: list[Platform] = [Platform.SWITCH, Platform.UPDATE]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -44,11 +45,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         ) from ex
 
     site_client = await client.get_site_client(OmadaSite(None, entry.data[CONF_SITE]))
-
+    controller = OmadaSiteController(hass, site_client)
-    hass.data[DOMAIN][entry.entry_id] = site_client
+    hass.data[DOMAIN][entry.entry_id] = controller
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
-
     return True
 
 

homeassistant/components/tplink_omada/controller.py

@@ -0,0 +1,52 @@
+"""Controller for sharing Omada API coordinators between platforms."""
+
+from functools import partial
+
+from tplink_omada_client.devices import OmadaSwitch, OmadaSwitchPortDetails
+from tplink_omada_client.omadasiteclient import OmadaSiteClient
+
+from homeassistant.core import HomeAssistant
+
+from .coordinator import OmadaCoordinator
+
+
+async def _poll_switch_state(
+    client: OmadaSiteClient, network_switch: OmadaSwitch
+) -> dict[str, OmadaSwitchPortDetails]:
+    """Poll a switch's current state."""
+    ports = await client.get_switch_ports(network_switch)
+    return {p.port_id: p for p in ports}
+
+
+class OmadaSiteController:
+    """Controller for the Omada SDN site."""
+
+    def __init__(self, hass: HomeAssistant, omada_client: OmadaSiteClient) -> None:
+        """Create the controller."""
+        self._hass = hass
+        self._omada_client = omada_client
+
+        self._switch_port_coordinators: dict[
+            str, OmadaCoordinator[OmadaSwitchPortDetails]
+        ] = {}
+
+    @property
+    def omada_client(self) -> OmadaSiteClient:
+        """Get the connected client API for the site to manage."""
+        return self._omada_client
+
+    def get_switch_port_coordinator(
+        self, switch: OmadaSwitch
+    ) -> OmadaCoordinator[OmadaSwitchPortDetails]:
+        """Get coordinator for network port information of a given switch."""
+        if switch.mac not in self._switch_port_coordinators:
+            self._switch_port_coordinators[switch.mac] = OmadaCoordinator[
+                OmadaSwitchPortDetails
+            ](
+                self._hass,
+                self._omada_client,
+                f"{switch.name} Ports",
+                partial(_poll_switch_state, network_switch=switch),
+            )
+
+        return self._switch_port_coordinators[switch.mac]

homeassistant/components/tplink_omada/coordinator.py

@@ -6,7 +6,7 @@ from typing import Generic, TypeVar
 
 import async_timeout
 from tplink_omada_client.exceptions import OmadaClientException
-from tplink_omada_client.omadaclient import OmadaClient
+from tplink_omada_client.omadaclient import OmadaSiteClient
 
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
@@ -22,15 +22,17 @@ class OmadaCoordinator(DataUpdateCoordinator[dict[str, T]], Generic[T]):
     def __init__(
         self,
         hass: HomeAssistant,
-        omada_client: OmadaClient,
+        omada_client: OmadaSiteClient,
-        update_func: Callable[[OmadaClient], Awaitable[dict[str, T]]],
+        name: str,
+        update_func: Callable[[OmadaSiteClient], Awaitable[dict[str, T]]],
+        poll_delay: int = 300,
     ) -> None:
         """Initialize my coordinator."""
         super().__init__(
             hass,
             _LOGGER,
-            name="Omada API Data",
+            name=f"Omada API Data - {name}",
-            update_interval=timedelta(seconds=300),
+            update_interval=timedelta(seconds=poll_delay),
         )
         self.omada_client = omada_client
         self._update_func = update_func

homeassistant/components/tplink_omada/entity.py

@@ -1,5 +1,7 @@
 """Base entity definitions."""
-from tplink_omada_client.devices import OmadaSwitch, OmadaSwitchPortDetails
+from typing import Generic, TypeVar
+
+from tplink_omada_client.devices import OmadaDevice
 
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
@@ -8,16 +10,14 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from .const import DOMAIN
 from .coordinator import OmadaCoordinator
 
+T = TypeVar("T")
 
-class OmadaSwitchDeviceEntity(
-    CoordinatorEntity[OmadaCoordinator[OmadaSwitchPortDetails]]
-):
-    """Common base class for all entities attached to Omada network switches."""
 
-    def __init__(
+class OmadaDeviceEntity(CoordinatorEntity[OmadaCoordinator[T]], Generic[T]):
-        self, coordinator: OmadaCoordinator[OmadaSwitchPortDetails], device: OmadaSwitch
+    """Common base class for all entities associated with Omada SDN Devices."""
-    ) -> None:
+
-        """Initialize the switch."""
+    def __init__(self, coordinator: OmadaCoordinator[T], device: OmadaDevice) -> None:
+        """Initialize the device."""
         super().__init__(coordinator)
         self.device = device
 

homeassistant/components/tplink_omada/manifest.json

@@ -6,5 +6,5 @@
   "documentation": "https://www.home-assistant.io/integrations/tplink_omada",
   "integration_type": "hub",
   "iot_class": "local_polling",
-  "requirements": ["tplink-omada-client==1.1.0"]
+  "requirements": ["tplink-omada-client==1.1.3"]
 }

homeassistant/components/tplink_omada/switch.py

@@ -1,12 +1,11 @@
 """Support for TPLink Omada device toggle options."""
 from __future__ import annotations
 
-from functools import partial
 from typing import Any
 
 from tplink_omada_client.definitions import PoEMode
 from tplink_omada_client.devices import OmadaSwitch, OmadaSwitchPortDetails
-from tplink_omada_client.omadasiteclient import OmadaSiteClient, SwitchPortOverrides
+from tplink_omada_client.omadasiteclient import SwitchPortOverrides
 
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.config_entries import ConfigEntry
@@ -15,27 +14,21 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
+from .controller import OmadaSiteController
 from .coordinator import OmadaCoordinator
-from .entity import OmadaSwitchDeviceEntity
+from .entity import OmadaDeviceEntity
 
 POE_SWITCH_ICON = "mdi:ethernet"
 
 
-async def poll_switch_state(
-    client: OmadaSiteClient, network_switch: OmadaSwitch
-) -> dict[str, OmadaSwitchPortDetails]:
-    """Poll a switch's current state."""
-    ports = await client.get_switch_ports(network_switch)
-    return {p.port_id: p for p in ports}
-
-
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up switches."""
-    omada_client: OmadaSiteClient = hass.data[DOMAIN][config_entry.entry_id]
+    controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
+    omada_client = controller.omada_client
 
     # Naming fun. Omada switches, as in the network hardware
     network_switches = await omada_client.get_switches()
@@ -44,10 +37,7 @@ async def async_setup_entry(
     for switch in [
         ns for ns in network_switches if ns.device_capabilities.supports_poe
     ]:
-        coordinator = OmadaCoordinator[OmadaSwitchPortDetails](
+        coordinator = controller.get_switch_port_coordinator(switch)
-            hass, omada_client, partial(poll_switch_state, network_switch=switch)
-        )
-
         await coordinator.async_request_refresh()
 
         for idx, port_id in enumerate(coordinator.data):
@@ -67,7 +57,9 @@ def get_port_base_name(port: OmadaSwitchPortDetails) -> str:
     return f"Port {port.port} ({port.name})"
 
 
-class OmadaNetworkSwitchPortPoEControl(OmadaSwitchDeviceEntity, SwitchEntity):
+class OmadaNetworkSwitchPortPoEControl(
+    OmadaDeviceEntity[OmadaSwitchPortDetails], SwitchEntity
+):
     """Representation of a PoE control toggle on a single network port on a switch."""
 
     _attr_has_entity_name = True

homeassistant/components/tplink_omada/update.py

@@ -0,0 +1,149 @@
+"""Support for TPLink Omada device toggle options."""
+from __future__ import annotations
+
+import logging
+from typing import Any, NamedTuple
+
+from tplink_omada_client.devices import OmadaFirmwareUpdate, OmadaListDevice
+from tplink_omada_client.omadasiteclient import OmadaSiteClient
+
+from homeassistant.components.update import UpdateEntity, UpdateEntityFeature
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.event import async_call_later
+
+from .const import DOMAIN
+from .controller import OmadaSiteController
+from .coordinator import OmadaCoordinator
+from .entity import OmadaDeviceEntity
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class FirmwareUpdateStatus(NamedTuple):
+    """Firmware update information for Omada SDN devices."""
+
+    device: OmadaListDevice
+    firmware: OmadaFirmwareUpdate | None
+
+
+async def _get_firmware_updates(client: OmadaSiteClient) -> list[FirmwareUpdateStatus]:
+    devices = await client.get_devices()
+    return [
+        FirmwareUpdateStatus(
+            device=d,
+            firmware=None
+            if not d.need_upgrade
+            else await client.get_firmware_details(d),
+        )
+        for d in devices
+    ]
+
+
+async def _poll_firmware_updates(
+    client: OmadaSiteClient,
+) -> dict[str, FirmwareUpdateStatus]:
+    """Poll the state of Omada Devices firmware update availability."""
+    return {d.device.mac: d for d in await _get_firmware_updates(client)}
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up switches."""
+    controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
+    omada_client = controller.omada_client
+
+    devices = await omada_client.get_devices()
+
+    coordinator = OmadaCoordinator[FirmwareUpdateStatus](
+        hass,
+        omada_client,
+        "Firmware Updates",
+        _poll_firmware_updates,
+        poll_delay=6 * 60 * 60,
+    )
+
+    entities: list = []
+    for device in devices:
+        entities.append(OmadaDeviceUpdate(coordinator, device))
+
+    async_add_entities(entities)
+    await coordinator.async_request_refresh()
+
+
+class OmadaDeviceUpdate(
+    OmadaDeviceEntity[FirmwareUpdateStatus],
+    UpdateEntity,
+):
+    """Firmware update status for Omada SDN devices."""
+
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL
+        | UpdateEntityFeature.PROGRESS
+        | UpdateEntityFeature.RELEASE_NOTES
+    )
+    _firmware_update: OmadaFirmwareUpdate = None
+
+    def __init__(
+        self,
+        coordinator: OmadaCoordinator[FirmwareUpdateStatus],
+        device: OmadaListDevice,
+    ) -> None:
+        """Initialize the update entity."""
+        super().__init__(coordinator, device)
+
+        self._mac = device.mac
+        self._device = device
+        self._omada_client = coordinator.omada_client
+
+        self._attr_unique_id = f"{device.mac}_firmware"
+        self._attr_has_entity_name = True
+        self._attr_name = "Firmware Update"
+        self._refresh_state()
+
+    def _refresh_state(self) -> None:
+        if self._firmware_update and self._device.need_upgrade:
+            self._attr_installed_version = self._firmware_update.current_version
+            self._attr_latest_version = self._firmware_update.latest_version
+        else:
+            self._attr_installed_version = self._device.firmware_version
+            self._attr_latest_version = self._device.firmware_version
+        self._attr_in_progress = self._device.fw_download
+
+        if self._attr_in_progress:
+            # While firmware update is in progress, poll more frequently
+            async_call_later(self.hass, 60, self._request_refresh)
+
+    async def _request_refresh(self, _now: Any) -> None:
+        await self.coordinator.async_request_refresh()
+
+    def release_notes(self) -> str | None:
+        """Get the release notes for the latest update."""
+        if self._firmware_update:
+            return str(self._firmware_update.release_notes)
+        return ""
+
+    async def async_install(
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Install a firmware update."""
+        if self._firmware_update and (
+            version is None or self._firmware_update.latest_version == version
+        ):
+            await self._omada_client.start_firmware_upgrade(self._device)
+            await self.coordinator.async_request_refresh()
+        else:
+            _LOGGER.error("Firmware upgrade is not available for %s", self._device.name)
+
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Handle updated data from the coordinator."""
+        status = self.coordinator.data[self._mac]
+        self._device = status.device
+        self._firmware_update = status.firmware
+        self._refresh_state()
+        self.async_write_ha_state()

requirements_all.txt

@@ -2524,7 +2524,7 @@ total_connect_client==2023.2
 tp-connected==0.0.4
 
 # homeassistant.components.tplink_omada
-tplink-omada-client==1.1.0
+tplink-omada-client==1.1.3
 
 # homeassistant.components.transmission
 transmission-rpc==3.4.0

requirements_test_all.txt

@@ -1785,7 +1785,7 @@ toonapi==0.2.1
 total_connect_client==2023.2
 
 # homeassistant.components.tplink_omada
-tplink-omada-client==1.1.0
+tplink-omada-client==1.1.3
 
 # homeassistant.components.transmission
 transmission-rpc==3.4.0