Fix unique identifiers where multiple IKEA Tradfri gateways are in use (#136060)

* Create unique identifiers where multiple gateways are in use

Resolving issue https://github.com/home-assistant/core/issues/134497

* Added migration function to __init__.py

Added migration function to execute upon initialisation, to:
a) remove the erroneously-added config)_entry added to the device (gateway B gets added as a config_entry to a device associated to gateway A), and
b) swap out the non-unique identifiers for genuinely unique identifiers.

* Added tests to simulate migration from bad data scenario (i.e. explicitly executing migrate_entity_unique_ids() from __init__.py)

* Ammendments suggested in first review

* Changes after second review

* Rewrite of test_migrate_config_entry_and_identifiers after feedback

* Converted migrate function into major version, updated tests

* Finalised variable naming convention per feedback, added test to validate config entry migrated to v2

* Hopefully final changes for cosmetic / comment stucture

* Further code-coverage in test_migrate_config_entry_and_identifiers()

* Minor test corrections

* Added test for non-tradfri identifiers

homeassistant/components/tradfri/__init__.py

@@ -159,7 +159,7 @@ def remove_stale_devices(
     device_entries = dr.async_entries_for_config_entry(
         device_registry, config_entry.entry_id
     )
-    all_device_ids = {device.id for device in devices}
+    all_device_ids = {str(device.id) for device in devices}
 
     for device_entry in device_entries:
         device_id: str | None = None
@@ -176,7 +176,7 @@ def remove_stale_devices(
                 gateway_id = _id
                 break
 
-            device_id = _id
+            device_id = _id.replace(f"{config_entry.data[CONF_GATEWAY_ID]}-", "")
             break
 
         if gateway_id is not None:
@@ -190,3 +190,93 @@ def remove_stale_devices(
             device_registry.async_update_device(
                 device_entry.id, remove_config_entry_id=config_entry.entry_id
             )
+
+
+async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+    """Migrate old entry."""
+    LOGGER.debug(
+        "Migrating Tradfri configuration from version %s.%s",
+        config_entry.version,
+        config_entry.minor_version,
+    )
+
+    if config_entry.version > 1:
+        # This means the user has downgraded from a future version
+        return False
+
+    if config_entry.version == 1:
+        # Migrate to version 2
+        migrate_config_entry_and_identifiers(hass, config_entry)
+
+        hass.config_entries.async_update_entry(config_entry, version=2)
+
+    LOGGER.debug(
+        "Migration to Tradfri configuration version %s.%s successful",
+        config_entry.version,
+        config_entry.minor_version,
+    )
+
+    return True
+
+
+def migrate_config_entry_and_identifiers(
+    hass: HomeAssistant, config_entry: ConfigEntry
+) -> None:
+    """Migrate old non-unique identifiers to new unique identifiers."""
+
+    related_device_flag: bool
+    device_id: str
+
+    device_reg = dr.async_get(hass)
+    # Get all devices associated to contextual gateway config_entry
+    # and loop through list of devices.
+    for device in dr.async_entries_for_config_entry(device_reg, config_entry.entry_id):
+        related_device_flag = False
+        for identifier in device.identifiers:
+            if identifier[0] != DOMAIN:
+                continue
+
+            related_device_flag = True
+
+            _id = identifier[1]
+
+            # Identify gateway device.
+            if _id == config_entry.data[CONF_GATEWAY_ID]:
+                # Using this to avoid updating gateway's own device registry entry
+                related_device_flag = False
+                break
+
+            device_id = str(_id)
+            break
+
+        # Check that device is related to tradfri domain (and is not the gateway itself)
+        if not related_device_flag:
+            continue
+
+        # Loop through list of config_entry_ids for device
+        config_entry_ids = device.config_entries
+        for config_entry_id in config_entry_ids:
+            # Check that the config entry in list is not the device's primary config entry
+            if config_entry_id == device.primary_config_entry:
+                continue
+
+            # Check that the 'other' config entry is also a tradfri config entry
+            other_entry = hass.config_entries.async_get_entry(config_entry_id)
+
+            if other_entry is None or other_entry.domain != DOMAIN:
+                continue
+
+            # Remove non-primary 'tradfri' config entry from device's config_entry_ids
+            device_reg.async_update_device(
+                device.id, remove_config_entry_id=config_entry_id
+            )
+
+        if config_entry.data[CONF_GATEWAY_ID] in device_id:
+            continue
+
+        device_reg.async_update_device(
+            device.id,
+            new_identifiers={
+                (DOMAIN, f"{config_entry.data[CONF_GATEWAY_ID]}-{device_id}")
+            },
+        )

homeassistant/components/tradfri/config_flow.py

@@ -35,7 +35,7 @@ class AuthError(Exception):
 class FlowHandler(ConfigFlow, domain=DOMAIN):
     """Handle a config flow."""
 
-    VERSION = 1
+    VERSION = 2
 
     def __init__(self) -> None:
         """Initialize flow."""

homeassistant/components/tradfri/entity.py

@@ -58,7 +58,7 @@ class TradfriBaseEntity(CoordinatorEntity[TradfriDeviceDataUpdateCoordinator]):
 
         info = self._device.device_info
         self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, self._device_id)},
+            identifiers={(DOMAIN, f"{gateway_id}-{self._device_id}")},
             manufacturer=info.manufacturer,
             model=info.model_number,
             name=self._device.name,

tests/components/tradfri/__init__.py

@@ -1,4 +1,6 @@
 """Tests for the tradfri component."""
 
 GATEWAY_ID = "mock-gateway-id"
+GATEWAY_ID1 = "mockgatewayid1"
+GATEWAY_ID2 = "mockgatewayid2"
 TRADFRI_PATH = "homeassistant.components.tradfri"

tests/components/tradfri/test_init.py

@@ -2,13 +2,19 @@
 
 from unittest.mock import MagicMock
 
+from pytradfri.const import ATTR_FIRMWARE_VERSION, ATTR_GATEWAY_ID
+from pytradfri.gateway import Gateway
+
 from homeassistant.components import tradfri
+from homeassistant.components.tradfri.const import DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
+from homeassistant.setup import async_setup_component
 
-from . import GATEWAY_ID
+from . import GATEWAY_ID, GATEWAY_ID1, GATEWAY_ID2
+from .common import CommandStore
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, load_json_object_fixture
 
 
 async def test_entry_setup_unload(
@@ -66,6 +72,7 @@ async def test_remove_stale_devices(
     device_registry.async_get_or_create(
         config_entry_id=config_entry.entry_id,
         identifiers={(tradfri.DOMAIN, "stale_device_id")},
+        name="stale-device",
     )
     device_entries = dr.async_entries_for_config_entry(
         device_registry, config_entry.entry_id
@@ -91,3 +98,178 @@ async def test_remove_stale_devices(
     assert device_entry.manufacturer == "IKEA of Sweden"
     assert device_entry.name == "Gateway"
     assert device_entry.model == "E1526"
+
+
+async def test_migrate_config_entry_and_identifiers(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    command_store: CommandStore,
+) -> None:
+    """Test correction of device registry entries."""
+    config_entry1 = MockConfigEntry(
+        domain=tradfri.DOMAIN,
+        data={
+            tradfri.CONF_HOST: "mock-host1",
+            tradfri.CONF_IDENTITY: "mock-identity1",
+            tradfri.CONF_KEY: "mock-key1",
+            tradfri.CONF_GATEWAY_ID: GATEWAY_ID1,
+        },
+    )
+
+    gateway1 = mock_gateway_fixture(command_store, GATEWAY_ID1)
+    command_store.register_device(
+        gateway1, load_json_object_fixture("bulb_w.json", DOMAIN)
+    )
+    config_entry1.add_to_hass(hass)
+
+    config_entry2 = MockConfigEntry(
+        domain=tradfri.DOMAIN,
+        data={
+            tradfri.CONF_HOST: "mock-host2",
+            tradfri.CONF_IDENTITY: "mock-identity2",
+            tradfri.CONF_KEY: "mock-key2",
+            tradfri.CONF_GATEWAY_ID: GATEWAY_ID2,
+        },
+    )
+
+    config_entry2.add_to_hass(hass)
+
+    # Add non-tradfri config entry for use in testing negation logic
+    config_entry3 = MockConfigEntry(
+        domain="test_domain",
+    )
+
+    config_entry3.add_to_hass(hass)
+
+    # Create gateway device for config entry 1
+    gateway1_device = device_registry.async_get_or_create(
+        config_entry_id=config_entry1.entry_id,
+        identifiers={(config_entry1.domain, config_entry1.data["gateway_id"])},
+        name="Gateway",
+    )
+
+    # Create bulb 1 on gateway 1 in Device Registry - this has the old identifiers format
+    gateway1_bulb1 = device_registry.async_get_or_create(
+        config_entry_id=config_entry1.entry_id,
+        identifiers={(tradfri.DOMAIN, 65537)},
+        name="bulb1",
+    )
+
+    # Update bulb 1 device to have both config entry IDs
+    # This is to simulate existing data scenario with older version of tradfri component
+    device_registry.async_update_device(
+        gateway1_bulb1.id,
+        add_config_entry_id=config_entry2.entry_id,
+    )
+
+    # Create bulb 2 on gateway 1 in Device Registry - this has the new identifiers format
+    gateway1_bulb2 = device_registry.async_get_or_create(
+        config_entry_id=config_entry1.entry_id,
+        identifiers={(tradfri.DOMAIN, f"{GATEWAY_ID1}-65538")},
+        name="bulb2",
+    )
+
+    # Update bulb 2 device to have an additional config entry from config_entry3
+    # This is to simulate scenario whereby a device entry
+    # is shared by multiple config entries
+    # and where at least one of those config entries is not the 'tradfri' domain
+    device_registry.async_update_device(
+        gateway1_bulb2.id,
+        add_config_entry_id=config_entry3.entry_id,
+        merge_identifiers={("test_domain", "config_entry_3-device2")},
+    )
+
+    # Create a device on config entry 3 in Device Registry
+    config_entry3_device = device_registry.async_get_or_create(
+        config_entry_id=config_entry3.entry_id,
+        identifiers={("test_domain", "config_entry_3-device1")},
+        name="device",
+    )
+
+    # Set up all tradfri config entries.
+    await async_setup_component(hass, DOMAIN, {})
+    await hass.async_block_till_done()
+
+    # Validate that gateway 1 bulb 1 is still the same device entry
+    # This inherently also validates that the device's identifiers
+    # have been updated to the new unique format
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry1.entry_id
+    )
+    assert (
+        device_registry.async_get_device(
+            identifiers={(tradfri.DOMAIN, f"{GATEWAY_ID1}-65537")}
+        ).id
+        == gateway1_bulb1.id
+    )
+
+    # Validate that gateway 1 bulb 1 only has gateway 1's config ID associated to it
+    # (Device at index 0 is the gateway)
+    assert device_entries[1].config_entries == {config_entry1.entry_id}
+
+    # Validate that the gateway 1 device is unchanged
+    assert device_entries[0].id == gateway1_device.id
+    assert device_entries[0].identifiers == gateway1_device.identifiers
+    assert device_entries[0].config_entries == gateway1_device.config_entries
+
+    # Validate that gateway 1 bulb 2 now only exists associated to config entry 3.
+    # The device will have had its identifiers updated to the new format (for the tradfri
+    # domain) per migrate_config_entry_and_identifiers().
+    # The device will have then been removed from config entry 1 (gateway1)
+    # due to it not matching a device in the command store.
+    device_entry = device_registry.async_get_device(
+        identifiers={(tradfri.DOMAIN, f"{GATEWAY_ID1}-65538")}
+    )
+
+    assert device_entry.id == gateway1_bulb2.id
+    # Assert that the only config entry associated to this device is config entry 3
+    assert device_entry.config_entries == {config_entry3.entry_id}
+    # Assert that that device's other identifiers remain untouched
+    assert device_entry.identifiers == {
+        (tradfri.DOMAIN, f"{GATEWAY_ID1}-65538"),
+        ("test_domain", "config_entry_3-device2"),
+    }
+
+    # Validate that gateway 2 bulb 1 has been added to device registry and with correct unique identifiers
+    # (This bulb device exists on gateway 2 because the command_store created above will be executed
+    # for each gateway being set up.)
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry2.entry_id
+    )
+    assert len(device_entries) == 2
+    assert device_entries[1].identifiers == {(tradfri.DOMAIN, f"{GATEWAY_ID2}-65537")}
+
+    # Validate that gateway 2 bulb 1 only has gateway 2's config ID associated to it
+    assert device_entries[1].config_entries == {config_entry2.entry_id}
+
+    # Validate that config entry 3 device 1 is still present,
+    # and has not had its config entries or identifiers changed
+    # N.B. The gateway1_bulb2 device will qualify in this set
+    # because the config entry 3 was added to it above
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry3.entry_id
+    )
+    assert len(device_entries) == 2
+    assert device_entries[0].id == config_entry3_device.id
+    assert device_entries[0].identifiers == {("test_domain", "config_entry_3-device1")}
+    assert device_entries[0].config_entries == {config_entry3.entry_id}
+
+    # Assert that the tradfri config entries have been migrated to v2 and
+    # the non-tradfri config entry remains at v1
+    assert config_entry1.version == 2
+    assert config_entry2.version == 2
+    assert config_entry3.version == 1
+
+
+def mock_gateway_fixture(command_store: CommandStore, gateway_id: str) -> Gateway:
+    """Mock a Tradfri gateway."""
+    gateway = Gateway()
+    command_store.register_response(
+        gateway.get_gateway_info(),
+        {ATTR_GATEWAY_ID: gateway_id, ATTR_FIRMWARE_VERSION: "1.2.1234"},
+    )
+    command_store.register_response(
+        gateway.get_devices(),
+        [],
+    )
+    return gateway