From 6271fe333d888c6145cd3478e4976160557445a4 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 24 Nov 2023 11:18:55 +0100 Subject: [PATCH] Rework some UniFi unique IDs (#104390) --- .../components/unifi/device_tracker.py | 32 +++++++- homeassistant/components/unifi/switch.py | 42 ++++++++--- tests/components/unifi/test_device_tracker.py | 9 ++- tests/components/unifi/test_switch.py | 75 ++++++++++++++++++- 4 files changed, 144 insertions(+), 14 deletions(-) diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 5c9694c669c..1be52b97974 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -17,14 +17,15 @@ from aiounifi.models.client import Client from aiounifi.models.device import Device from aiounifi.models.event import Event, EventKey -from homeassistant.components.device_tracker import ScannerEntity, SourceType +from homeassistant.components.device_tracker import DOMAIN, ScannerEntity, SourceType from homeassistant.config_entries import ConfigEntry from homeassistant.core import Event as core_Event, HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddEntitiesCallback +import homeassistant.helpers.entity_registry as er import homeassistant.util.dt as dt_util -from .controller import UniFiController +from .controller import UNIFI_DOMAIN, UniFiController from .entity import ( HandlerT, UnifiEntity, @@ -175,7 +176,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = ( object_fn=lambda api, obj_id: api.clients[obj_id], should_poll=False, supported_fn=lambda controller, obj_id: True, - unique_id_fn=lambda controller, obj_id: f"{obj_id}-{controller.site}", + unique_id_fn=lambda controller, obj_id: f"{controller.site}-{obj_id}", ip_address_fn=lambda api, obj_id: api.clients[obj_id].ip, hostname_fn=lambda api, obj_id: api.clients[obj_id].hostname, ), @@ -201,12 +202,37 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = ( ) +@callback +def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> None: + """Normalize client unique ID to have a prefix rather than suffix. + + Introduced with release 2023.12. + """ + controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + ent_reg = er.async_get(hass) + + @callback + def update_unique_id(obj_id: str) -> None: + """Rework unique ID.""" + new_unique_id = f"{controller.site}-{obj_id}" + if ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, new_unique_id): + return + + unique_id = f"{obj_id}-{controller.site}" + if entity_id := ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, unique_id): + ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id) + + for obj_id in list(controller.api.clients) + list(controller.api.clients_all): + update_unique_id(obj_id) + + async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up device tracker for UniFi Network integration.""" + async_update_unique_id(hass, config_entry) UniFiController.register_platform( hass, config_entry, async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS ) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 18a3dbc3b90..1e9ec8b14c8 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -42,9 +42,10 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback +import homeassistant.helpers.entity_registry as er from .const import ATTR_MANUFACTURER -from .controller import UniFiController +from .controller import UNIFI_DOMAIN, UniFiController from .entity import ( HandlerT, SubscriptionT, @@ -199,12 +200,6 @@ class UnifiSwitchEntityDescription( only_event_for_state_change: bool = False -def _make_unique_id(obj_id: str, type_name: str) -> str: - """Split an object id by the first underscore and interpose the given type.""" - prefix, _, suffix = obj_id.partition("_") - return f"{prefix}-{type_name}-{suffix}" - - ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( UnifiSwitchEntityDescription[Clients, Client]( key="Block client", @@ -262,7 +257,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( object_fn=lambda api, obj_id: api.outlets[obj_id], should_poll=False, supported_fn=async_outlet_supports_switching_fn, - unique_id_fn=lambda controller, obj_id: _make_unique_id(obj_id, "outlet"), + unique_id_fn=lambda controller, obj_id: f"outlet-{obj_id}", ), UnifiSwitchEntityDescription[PortForwarding, PortForward]( key="Port forward control", @@ -303,7 +298,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( object_fn=lambda api, obj_id: api.ports[obj_id], should_poll=False, supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe, - unique_id_fn=lambda controller, obj_id: _make_unique_id(obj_id, "poe"), + unique_id_fn=lambda controller, obj_id: f"poe-{obj_id}", ), UnifiSwitchEntityDescription[Wlans, Wlan]( key="WLAN control", @@ -328,12 +323,41 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( ) +@callback +def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> None: + """Normalize switch unique ID to have a prefix rather than midfix. + + Introduced with release 2023.12. + """ + controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + ent_reg = er.async_get(hass) + + @callback + def update_unique_id(obj_id: str, type_name: str) -> None: + """Rework unique ID.""" + new_unique_id = f"{type_name}-{obj_id}" + if ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, new_unique_id): + return + + prefix, _, suffix = obj_id.partition("_") + unique_id = f"{prefix}-{type_name}-{suffix}" + if entity_id := ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, unique_id): + ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id) + + for obj_id in controller.api.outlets: + update_unique_id(obj_id, "outlet") + + for obj_id in controller.api.ports: + update_unique_id(obj_id, "poe") + + async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up switches for UniFi Network integration.""" + async_update_unique_id(hass, config_entry) UniFiController.register_platform( hass, config_entry, diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index cbff868d9a6..abe12a1e243 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -939,13 +939,20 @@ async def test_restoring_client( ) registry = er.async_get(hass) - registry.async_get_or_create( + registry.async_get_or_create( # Unique ID updated TRACKER_DOMAIN, UNIFI_DOMAIN, f'{restored["mac"]}-site_id', suggested_object_id=restored["hostname"], config_entry=config_entry, ) + registry.async_get_or_create( # Unique ID already updated + TRACKER_DOMAIN, + UNIFI_DOMAIN, + f'site_id-{client["mac"]}', + suggested_object_id=client["hostname"], + config_entry=config_entry, + ) await setup_unifi_integration( hass, diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index fe2ee5dc9e8..00ebcd0e683 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -5,6 +5,7 @@ from datetime import timedelta from aiounifi.models.message import MessageKey import pytest +from homeassistant import config_entries from homeassistant.components.switch import ( DOMAIN as SWITCH_DOMAIN, SERVICE_TURN_OFF, @@ -32,7 +33,12 @@ from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_registry import RegistryEntryDisabler from homeassistant.util import dt as dt_util -from .test_controller import CONTROLLER_HOST, SITE, setup_unifi_integration +from .test_controller import ( + CONTROLLER_HOST, + ENTRY_CONFIG, + SITE, + setup_unifi_integration, +) from tests.common import async_fire_time_changed from tests.test_util.aiohttp import AiohttpClientMocker @@ -1585,3 +1591,70 @@ async def test_port_forwarding_switches( mock_unifi_websocket(message=MessageKey.PORT_FORWARD_DELETED, data=_data) await hass.async_block_till_done() assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 + + +async def test_updating_unique_id( + hass: HomeAssistant, aioclient_mock: AiohttpClientMocker +) -> None: + """Verify outlet control and poe control unique ID update works.""" + poe_device = { + "board_rev": 3, + "device_id": "mock-id", + "ip": "10.0.0.1", + "last_seen": 1562600145, + "mac": "00:00:00:00:01:01", + "model": "US16P150", + "name": "switch", + "state": 1, + "type": "usw", + "version": "4.0.42.10433", + "port_table": [ + { + "media": "GE", + "name": "Port 1", + "port_idx": 1, + "poe_caps": 7, + "poe_class": "Class 4", + "poe_enable": True, + "poe_mode": "auto", + "poe_power": "2.56", + "poe_voltage": "53.40", + "portconf_id": "1a1", + "port_poe": True, + "up": True, + }, + ], + } + + config_entry = config_entries.ConfigEntry( + version=1, + domain=UNIFI_DOMAIN, + title="Mock Title", + data=ENTRY_CONFIG, + source="test", + options={}, + entry_id="1", + ) + + registry = er.async_get(hass) + registry.async_get_or_create( + SWITCH_DOMAIN, + UNIFI_DOMAIN, + f'{poe_device["mac"]}-poe-1', + suggested_object_id="switch_port_1_poe", + config_entry=config_entry, + ) + registry.async_get_or_create( + SWITCH_DOMAIN, + UNIFI_DOMAIN, + f'{OUTLET_UP1["mac"]}-outlet-1', + suggested_object_id="plug_outlet_1", + config_entry=config_entry, + ) + + await setup_unifi_integration( + hass, aioclient_mock, devices_response=[poe_device, OUTLET_UP1] + ) + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2 + assert hass.states.get("switch.switch_port_1_poe") + assert hass.states.get("switch.plug_outlet_1")