TP-Link Omada update entities code review feedback (#89668)

homeassistant/components/tplink_omada/__init__.py

@@ -44,7 +44,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             f"Unexpected error connecting to Omada controller: {ex}"
         ) from ex
 
-    site_client = await client.get_site_client(OmadaSite(None, entry.data[CONF_SITE]))
+    site_client = await client.get_site_client(OmadaSite("", entry.data[CONF_SITE]))
     controller = OmadaSiteController(hass, site_client)
     hass.data[DOMAIN][entry.entry_id] = controller
 

homeassistant/components/tplink_omada/controller.py

@@ -1,7 +1,5 @@
 """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
 
@@ -9,13 +7,28 @@ from homeassistant.core import HomeAssistant
 
 from .coordinator import OmadaCoordinator
 
+POLL_SWITCH_PORT = 300
 
-async def _poll_switch_state(
+
-    client: OmadaSiteClient, network_switch: OmadaSwitch
+class OmadaSwitchPortCoordinator(OmadaCoordinator[OmadaSwitchPortDetails]):
-) -> dict[str, OmadaSwitchPortDetails]:
+    """Coordinator for getting details about ports on a switch."""
-    """Poll a switch's current state."""
+
-    ports = await client.get_switch_ports(network_switch)
+    def __init__(
-    return {p.port_id: p for p in ports}
+        self,
+        hass: HomeAssistant,
+        omada_client: OmadaSiteClient,
+        network_switch: OmadaSwitch,
+    ) -> None:
+        """Initialize my coordinator."""
+        super().__init__(
+            hass, omada_client, f"{network_switch.name} Ports", POLL_SWITCH_PORT
+        )
+        self._network_switch = network_switch
+
+    async def poll_update(self) -> dict[str, OmadaSwitchPortDetails]:
+        """Poll a switch's current state."""
+        ports = await self.omada_client.get_switch_ports(self._network_switch)
+        return {p.port_id: p for p in ports}
 
 
 class OmadaSiteController:
@@ -26,9 +39,7 @@ class OmadaSiteController:
         self._hass = hass
         self._omada_client = omada_client
 
-        self._switch_port_coordinators: dict[
+        self._switch_port_coordinators: dict[str, OmadaSwitchPortCoordinator] = {}
-            str, OmadaCoordinator[OmadaSwitchPortDetails]
-        ] = {}
 
     @property
     def omada_client(self) -> OmadaSiteClient:
@@ -37,16 +48,11 @@ class OmadaSiteController:
 
     def get_switch_port_coordinator(
         self, switch: OmadaSwitch
-    ) -> OmadaCoordinator[OmadaSwitchPortDetails]:
+    ) -> OmadaSwitchPortCoordinator:
         """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[
+            self._switch_port_coordinators[switch.mac] = OmadaSwitchPortCoordinator(
-                OmadaSwitchPortDetails
+                self._hass, self._omada_client, switch
-            ](
-                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

@@ -1,5 +1,4 @@
 """Generic Omada API coordinator."""
-from collections.abc import Awaitable, Callable
 from datetime import timedelta
 import logging
 from typing import Generic, TypeVar
@@ -24,7 +23,6 @@ class OmadaCoordinator(DataUpdateCoordinator[dict[str, T]], Generic[T]):
         hass: HomeAssistant,
         omada_client: OmadaSiteClient,
         name: str,
-        update_func: Callable[[OmadaSiteClient], Awaitable[dict[str, T]]],
         poll_delay: int = 300,
     ) -> None:
         """Initialize my coordinator."""
@@ -35,12 +33,15 @@ class OmadaCoordinator(DataUpdateCoordinator[dict[str, T]], Generic[T]):
             update_interval=timedelta(seconds=poll_delay),
         )
         self.omada_client = omada_client
-        self._update_func = update_func
 
     async def _async_update_data(self) -> dict[str, T]:
         """Fetch data from API endpoint."""
         try:
             async with async_timeout.timeout(10):
-                return await self._update_func(self.omada_client)
+                return await self.poll_update()
         except OmadaClientException as err:
             raise UpdateFailed(f"Error communicating with API: {err}") from err
+
+    async def poll_update(self) -> dict[str, T]:
+        """Poll the current data from the controller."""
+        raise NotImplementedError("Update method not implemented")

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.3"]
+  "requirements": ["tplink-omada-client==1.1.4"]
 }

homeassistant/components/tplink_omada/switch.py

@@ -14,8 +14,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .controller import OmadaSiteController
+from .controller import OmadaSiteController, OmadaSwitchPortCoordinator
-from .coordinator import OmadaCoordinator
 from .entity import OmadaDeviceEntity
 
 POE_SWITCH_ICON = "mdi:ethernet"
@@ -68,7 +67,7 @@ class OmadaNetworkSwitchPortPoEControl(
 
     def __init__(
         self,
-        coordinator: OmadaCoordinator[OmadaSwitchPortDetails],
+        coordinator: OmadaSwitchPortCoordinator,
         device: OmadaSwitch,
         port_id: str,
     ) -> None:

homeassistant/components/tplink_omada/update.py

@@ -1,24 +1,26 @@
-"""Support for TPLink Omada device toggle options."""
+"""Support for TPLink Omada device firmware updates."""
 from __future__ import annotations
 
-import logging
+from datetime import timedelta
 from typing import Any, NamedTuple
 
 from tplink_omada_client.devices import OmadaFirmwareUpdate, OmadaListDevice
+from tplink_omada_client.exceptions import OmadaClientException, RequestFailed
 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.exceptions import HomeAssistantError
 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__)
+POLL_DELAY_IDLE = 6 * 60 * 60
+POLL_DELAY_UPGRADE = 60
 
 
 class FirmwareUpdateStatus(NamedTuple):
@@ -28,24 +30,39 @@ class FirmwareUpdateStatus(NamedTuple):
     firmware: OmadaFirmwareUpdate | None
 
 
-async def _get_firmware_updates(client: OmadaSiteClient) -> list[FirmwareUpdateStatus]:
+class OmadaFirmwareUpdateCoodinator(OmadaCoordinator[FirmwareUpdateStatus]):
-    devices = await client.get_devices()
+    """Coordinator for getting details about ports on a switch."""
-    return [
+
-        FirmwareUpdateStatus(
+    def __init__(self, hass: HomeAssistant, omada_client: OmadaSiteClient) -> None:
-            device=d,
+        """Initialize my coordinator."""
-            firmware=None
+        super().__init__(hass, omada_client, "Firmware Updates", POLL_DELAY_IDLE)
-            if not d.need_upgrade
+
-            else await client.get_firmware_details(d),
+    async def _get_firmware_updates(self) -> list[FirmwareUpdateStatus]:
+        devices = await self.omada_client.get_devices()
+
+        updates = [
+            FirmwareUpdateStatus(
+                device=d,
+                firmware=None
+                if not d.need_upgrade
+                else await self.omada_client.get_firmware_details(d),
+            )
+            for d in devices
+        ]
+
+        # During a firmware upgrade, poll more frequently
+        self.update_interval = timedelta(
+            seconds=(
+                POLL_DELAY_UPGRADE
+                if any(u.device.fw_download for u in updates)
+                else POLL_DELAY_IDLE
+            )
         )
-        for d in devices
+        return updates
-    ]
 
-
+    async def poll_update(self) -> dict[str, FirmwareUpdateStatus]:
-async def _poll_firmware_updates(
+        """Poll the state of Omada Devices firmware update availability."""
-    client: OmadaSiteClient,
+        return {d.device.mac: d for d in await self._get_firmware_updates()}
-) -> 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(
@@ -59,19 +76,9 @@ async def async_setup_entry(
 
     devices = await omada_client.get_devices()
 
-    coordinator = OmadaCoordinator[FirmwareUpdateStatus](
+    coordinator = OmadaFirmwareUpdateCoodinator(hass, omada_client)
-        hass,
-        omada_client,
-        "Firmware Updates",
-        _poll_firmware_updates,
-        poll_delay=6 * 60 * 60,
-    )
 
-    entities: list = []
+    async_add_entities(OmadaDeviceUpdate(coordinator, device) for device in devices)
-    for device in devices:
-        entities.append(OmadaDeviceUpdate(coordinator, device))
-
-    async_add_entities(entities)
     await coordinator.async_request_refresh()
 
 
@@ -86,64 +93,57 @@ class OmadaDeviceUpdate(
         | UpdateEntityFeature.PROGRESS
         | UpdateEntityFeature.RELEASE_NOTES
     )
-    _firmware_update: OmadaFirmwareUpdate = None
+    _attr_has_entity_name = True
+    _attr_name = "Firmware update"
 
     def __init__(
         self,
-        coordinator: OmadaCoordinator[FirmwareUpdateStatus],
+        coordinator: OmadaFirmwareUpdateCoodinator,
         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:
+        status = self.coordinator.data[self._mac]
-            return str(self._firmware_update.release_notes)
+        if status.firmware:
-        return ""
+            return status.firmware.release_notes
+        return None
 
     async def async_install(
         self, version: str | None, backup: bool, **kwargs: Any
     ) -> None:
         """Install a firmware update."""
-        if self._firmware_update and (
+        try:
-            version is None or self._firmware_update.latest_version == version
+            await self._omada_client.start_firmware_upgrade(
-        ):
+                self.coordinator.data[self._mac].device
-            await self._omada_client.start_firmware_upgrade(self._device)
+            )
+        except RequestFailed as ex:
+            raise HomeAssistantError("Firmware update request rejected") from ex
+        except OmadaClientException as ex:
+            raise HomeAssistantError(
+                "Unable to send Firmware update request. Check the controller is online."
+            ) from ex
+        finally:
             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
+        if status.firmware and status.device.need_upgrade:
-        self._refresh_state()
+            self._attr_installed_version = status.firmware.current_version
+            self._attr_latest_version = status.firmware.latest_version
+        else:
+            self._attr_installed_version = status.device.firmware_version
+            self._attr_latest_version = status.device.firmware_version
+        self._attr_in_progress = status.device.fw_download
+
         self.async_write_ha_state()

requirements_all.txt

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

requirements_test_all.txt

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