From 41b4c5532de0a59de32ff87c8d8179f4732dbb4f Mon Sep 17 00:00:00 2001 From: MarkGodwin Date: Mon, 13 Mar 2023 01:26:34 +0000 Subject: [PATCH] Add Update entities to TP-Link Omada integration (#89562) * Bump tplink-omada * Add omada firmware updates * Excluded from code coverage * Fixed entity name --- .coveragerc | 2 + .../components/tplink_omada/__init__.py | 8 +- .../components/tplink_omada/controller.py | 52 ++++++ .../components/tplink_omada/coordinator.py | 12 +- .../components/tplink_omada/entity.py | 18 +-- .../components/tplink_omada/manifest.json | 2 +- .../components/tplink_omada/switch.py | 26 ++- .../components/tplink_omada/update.py | 149 ++++++++++++++++++ requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- 10 files changed, 235 insertions(+), 38 deletions(-) create mode 100644 homeassistant/components/tplink_omada/controller.py create mode 100644 homeassistant/components/tplink_omada/update.py diff --git a/.coveragerc b/.coveragerc index a533343bf06..20ee077ffa0 100644 --- a/.coveragerc +++ b/.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 diff --git a/homeassistant/components/tplink_omada/__init__.py b/homeassistant/components/tplink_omada/__init__.py index 1e7db69cc95..709ad520125 100644 --- a/homeassistant/components/tplink_omada/__init__.py +++ b/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])) - - hass.data[DOMAIN][entry.entry_id] = site_client + controller = OmadaSiteController(hass, site_client) + hass.data[DOMAIN][entry.entry_id] = controller await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) - return True diff --git a/homeassistant/components/tplink_omada/controller.py b/homeassistant/components/tplink_omada/controller.py new file mode 100644 index 00000000000..b42cb37ff76 --- /dev/null +++ b/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] diff --git a/homeassistant/components/tplink_omada/coordinator.py b/homeassistant/components/tplink_omada/coordinator.py index 6950e3b6d74..d73461dc786 100644 --- a/homeassistant/components/tplink_omada/coordinator.py +++ b/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, - update_func: Callable[[OmadaClient], Awaitable[dict[str, T]]], + omada_client: OmadaSiteClient, + 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", - update_interval=timedelta(seconds=300), + name=f"Omada API Data - {name}", + update_interval=timedelta(seconds=poll_delay), ) self.omada_client = omada_client self._update_func = update_func diff --git a/homeassistant/components/tplink_omada/entity.py b/homeassistant/components/tplink_omada/entity.py index c3cc1433b9c..41cb1c69180 100644 --- a/homeassistant/components/tplink_omada/entity.py +++ b/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__( - self, coordinator: OmadaCoordinator[OmadaSwitchPortDetails], device: OmadaSwitch - ) -> None: - """Initialize the switch.""" +class OmadaDeviceEntity(CoordinatorEntity[OmadaCoordinator[T]], Generic[T]): + """Common base class for all entities associated with Omada SDN Devices.""" + + def __init__(self, coordinator: OmadaCoordinator[T], device: OmadaDevice) -> None: + """Initialize the device.""" super().__init__(coordinator) self.device = device diff --git a/homeassistant/components/tplink_omada/manifest.json b/homeassistant/components/tplink_omada/manifest.json index 005589a2f99..a0fb58b3f6c 100644 --- a/homeassistant/components/tplink_omada/manifest.json +++ b/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"] } diff --git a/homeassistant/components/tplink_omada/switch.py b/homeassistant/components/tplink_omada/switch.py index dd5ee3168d2..e85b1c181fc 100644 --- a/homeassistant/components/tplink_omada/switch.py +++ b/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]( - hass, omada_client, partial(poll_switch_state, network_switch=switch) - ) - + coordinator = controller.get_switch_port_coordinator(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 diff --git a/homeassistant/components/tplink_omada/update.py b/homeassistant/components/tplink_omada/update.py new file mode 100644 index 00000000000..5581f61d824 --- /dev/null +++ b/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() diff --git a/requirements_all.txt b/requirements_all.txt index 742344b316b..a3f91219b3e 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index d631b0ecda2..ae5a35f3184 100644 --- a/requirements_test_all.txt +++ b/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