Change device identifier and binary_sensor unique_id for airOS (home-assistant#153085)

Co-authored-by: G Johansson <[email protected]>

Author: CoMPaTech

homeassistant/components/airos/__init__.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import logging
+
 from airos.airos8 import AirOS8
 
 from homeassistant.const import (
@@ -12,17 +14,20 @@
     CONF_VERIFY_SSL,
     Platform,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DEFAULT_SSL, DEFAULT_VERIFY_SSL, SECTION_ADVANCED_SETTINGS
+from .const import DEFAULT_SSL, DEFAULT_VERIFY_SSL, DOMAIN, SECTION_ADVANCED_SETTINGS
 from .coordinator import AirOSConfigEntry, AirOSDataUpdateCoordinator
 
 _PLATFORMS: list[Platform] = [
     Platform.BINARY_SENSOR,
     Platform.SENSOR,
 ]
 
+_LOGGER = logging.getLogger(__name__)
+
 
 async def async_setup_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> bool:
     """Set up Ubiquiti airOS from a config entry."""
@@ -54,11 +59,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> boo
 async def async_migrate_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> bool:
     """Migrate old config entry."""
 
-    if entry.version > 1:
-        # This means the user has downgraded from a future version
+    # This means the user has downgraded from a future version
+    if entry.version > 2:
         return False
 
+    # 1.1 Migrate config_entry to add advanced ssl settings
     if entry.version == 1 and entry.minor_version == 1:
+        new_minor_version = 2
         new_data = {**entry.data}
         advanced_data = {
             CONF_SSL: DEFAULT_SSL,
@@ -69,7 +76,52 @@ async def async_migrate_entry(hass: HomeAssistant, entry: AirOSConfigEntry) -> b
         hass.config_entries.async_update_entry(
             entry,
             data=new_data,
-            minor_version=2,
+            minor_version=new_minor_version,
+        )
+
+    # 2.1 Migrate binary_sensor entity unique_id from device_id to mac_address
+    #     Step 1 - migrate binary_sensor entity unique_id
+    #     Step 2 - migrate device entity identifier
+    if entry.version == 1:
+        new_version = 2
+        new_minor_version = 1
+
+        mac_adress = dr.format_mac(entry.unique_id)
+
+        device_registry = dr.async_get(hass)
+        if device_entry := device_registry.async_get_device(
+            connections={(dr.CONNECTION_NETWORK_MAC, mac_adress)}
+        ):
+            old_device_id = next(
+                (
+                    device_id
+                    for domain, device_id in device_entry.identifiers
+                    if domain == DOMAIN
+                ),
+            )
+
+            @callback
+            def update_unique_id(
+                entity_entry: er.RegistryEntry,
+            ) -> dict[str, str] | None:
+                """Update unique id from device_id to mac address."""
+                if old_device_id and entity_entry.unique_id.startswith(old_device_id):
+                    suffix = entity_entry.unique_id.removeprefix(old_device_id)
+                    new_unique_id = f"{mac_adress}{suffix}"
+                    return {"new_unique_id": new_unique_id}
+                return None
+
+            await er.async_migrate_entries(hass, entry.entry_id, update_unique_id)
+
+            new_identifiers = device_entry.identifiers.copy()
+            new_identifiers.discard((DOMAIN, old_device_id))
+            new_identifiers.add((DOMAIN, mac_adress))
+            device_registry.async_update_device(
+                device_entry.id, new_identifiers=new_identifiers
+            )
+
+        hass.config_entries.async_update_entry(
+            entry, version=new_version, minor_version=new_minor_version
         )
 
     return True

homeassistant/components/airos/binary_sensor.py

@@ -98,7 +98,7 @@ def __init__(
         super().__init__(coordinator)
 
         self.entity_description = description
-        self._attr_unique_id = f"{coordinator.data.host.device_id}_{description.key}"
+        self._attr_unique_id = f"{coordinator.data.derived.mac}_{description.key}"
 
     @property
     def is_on(self) -> bool:

homeassistant/components/airos/config_flow.py

@@ -57,8 +57,8 @@
 class AirOSConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Ubiquiti airOS."""
 
-    VERSION = 1
-    MINOR_VERSION = 2
+    VERSION = 2
+    MINOR_VERSION = 1
 
     def __init__(self) -> None:
         """Initialize the config flow."""

homeassistant/components/airos/entity.py

@@ -33,7 +33,7 @@ def __init__(self, coordinator: AirOSDataUpdateCoordinator) -> None:
         self._attr_device_info = DeviceInfo(
             connections={(CONNECTION_NETWORK_MAC, airos_data.derived.mac)},
             configuration_url=configuration_url,
-            identifiers={(DOMAIN, str(airos_data.host.device_id))},
+            identifiers={(DOMAIN, airos_data.derived.mac)},
             manufacturer=MANUFACTURER,
             model=airos_data.host.devmodel,
             model_id=(

tests/components/airos/snapshots/test_binary_sensor.ambr

@@ -30,7 +30,7 @@
     'suggested_object_id': None,
     'supported_features': 0,
     'translation_key': 'dhcp_client',
-    'unique_id': '03aa0d0b40fed0a47088293584ef5432_dhcp_client',
+    'unique_id': '01:23:45:67:89:AB_dhcp_client',
     'unit_of_measurement': None,
   })
 # ---
@@ -79,7 +79,7 @@
     'suggested_object_id': None,
     'supported_features': 0,
     'translation_key': 'dhcp_server',
-    'unique_id': '03aa0d0b40fed0a47088293584ef5432_dhcp_server',
+    'unique_id': '01:23:45:67:89:AB_dhcp_server',
     'unit_of_measurement': None,
   })
 # ---
@@ -128,7 +128,7 @@
     'suggested_object_id': None,
     'supported_features': 0,
     'translation_key': 'dhcp6_server',
-    'unique_id': '03aa0d0b40fed0a47088293584ef5432_dhcp6_server',
+    'unique_id': '01:23:45:67:89:AB_dhcp6_server',
     'unit_of_measurement': None,
   })
 # ---
@@ -177,7 +177,7 @@
     'suggested_object_id': None,
     'supported_features': 0,
     'translation_key': 'port_forwarding',
-    'unique_id': '03aa0d0b40fed0a47088293584ef5432_portfw',
+    'unique_id': '01:23:45:67:89:AB_portfw',
     'unit_of_measurement': None,
   })
 # ---
@@ -225,7 +225,7 @@
     'suggested_object_id': None,
     'supported_features': 0,
     'translation_key': 'pppoe',
-    'unique_id': '03aa0d0b40fed0a47088293584ef5432_pppoe',
+    'unique_id': '01:23:45:67:89:AB_pppoe',
     'unit_of_measurement': None,
   })
 # ---

tests/components/airos/test_init.py

@@ -4,12 +4,16 @@
 
 from unittest.mock import ANY, MagicMock
 
+import pytest
+
 from homeassistant.components.airos.const import (
     DEFAULT_SSL,
     DEFAULT_VERIFY_SSL,
     DOMAIN,
     SECTION_ADVANCED_SETTINGS,
 )
+from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR_DOMAIN
+from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.config_entries import SOURCE_USER, ConfigEntryState
 from homeassistant.const import (
     CONF_HOST,
@@ -19,6 +23,7 @@
     CONF_VERIFY_SSL,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from tests.common import MockConfigEntry
 
@@ -108,8 +113,10 @@ async def test_setup_entry_without_ssl(
     assert entry.data[SECTION_ADVANCED_SETTINGS][CONF_VERIFY_SSL] is False
 
 
-async def test_migrate_entry(hass: HomeAssistant, mock_airos_client: MagicMock) -> None:
-    """Test migrate entry unique id."""
+async def test_ssl_migrate_entry(
+    hass: HomeAssistant, mock_airos_client: MagicMock
+) -> None:
+    """Test migrate entry SSL options."""
     entry = MockConfigEntry(
         domain=DOMAIN,
         source=SOURCE_USER,
@@ -124,11 +131,77 @@ async def test_migrate_entry(hass: HomeAssistant, mock_airos_client: MagicMock)
     await hass.async_block_till_done()
 
     assert entry.state is ConfigEntryState.LOADED
-    assert entry.version == 1
-    assert entry.minor_version == 2
+    assert entry.version == 2
+    assert entry.minor_version == 1
     assert entry.data == MOCK_CONFIG_V1_2
 
 
+@pytest.mark.parametrize(
+    ("sensor_domain", "sensor_name", "mock_id"),
+    [
+        (BINARY_SENSOR_DOMAIN, "port_forwarding", "device_id_12345"),
+        (SENSOR_DOMAIN, "antenna_gain", "01:23:45:67:89:ab"),
+    ],
+)
+async def test_uid_migrate_entry(
+    hass: HomeAssistant,
+    mock_airos_client: MagicMock,
+    device_registry: dr.DeviceRegistry,
+    sensor_domain: str,
+    sensor_name: str,
+    mock_id: str,
+) -> None:
+    """Test migrate entry unique id."""
+    entity_registry = er.async_get(hass)
+
+    MOCK_MAC = dr.format_mac("01:23:45:67:89:AB")
+    MOCK_ID = "device_id_12345"
+    old_unique_id = f"{mock_id}_{sensor_name}"
+    new_unique_id = f"{MOCK_MAC}_{sensor_name}"
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data=MOCK_CONFIG_V1_2,
+        entry_id="1",
+        unique_id=mock_id,
+        version=1,
+        minor_version=2,
+    )
+    entry.add_to_hass(hass)
+
+    device_registry.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={(DOMAIN, MOCK_ID)},
+        connections={
+            (dr.CONNECTION_NETWORK_MAC, MOCK_MAC),
+        },
+    )
+    await hass.async_block_till_done()
+
+    old_entity_entry = entity_registry.async_get_or_create(
+        DOMAIN, sensor_domain, old_unique_id, config_entry=entry
+    )
+    original_entity_id = old_entity_entry.entity_id
+
+    hass.config_entries.async_update_entry(entry, unique_id=MOCK_MAC)
+    await hass.async_block_till_done()
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    updated_entity_entry = entity_registry.async_get(original_entity_id)
+
+    assert entry.state is ConfigEntryState.LOADED
+    assert entry.version == 2
+    assert entry.minor_version == 1
+    assert (
+        entity_registry.async_get_entity_id(sensor_domain, DOMAIN, old_unique_id)
+        is None
+    )
+    assert updated_entity_entry.unique_id == new_unique_id
+
+
 async def test_migrate_future_return(
     hass: HomeAssistant,
     mock_airos_client: MagicMock,
@@ -140,7 +213,7 @@ async def test_migrate_future_return(
         data=MOCK_CONFIG_V1_2,
         entry_id="1",
         unique_id="airos_device",
-        version=2,
+        version=3,
     )
     entry.add_to_hass(hass)