From d0412d65ac0819975ab0ab3d0889c443c7964726 Mon Sep 17 00:00:00 2001
From: Martin Hjelmare <marhje52@gmail.com>
Date: Tue, 1 Feb 2022 01:43:16 +0100
Subject: [PATCH] Remove stale tradfri devices (#65218)

---
 homeassistant/components/tradfri/__init__.py | 47 +++++++++++++++++++-
 tests/components/tradfri/test_init.py        | 41 +++++++++++++++++
 2 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/tradfri/__init__.py b/homeassistant/components/tradfri/__init__.py
index 61e2b50d7aa..c11b9874bae 100644
--- a/homeassistant/components/tradfri/__init__.py
+++ b/homeassistant/components/tradfri/__init__.py
@@ -15,9 +15,10 @@ import voluptuous as vol
 from homeassistant import config_entries
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import Event, HomeAssistant
+from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
 import homeassistant.helpers.config_validation as cv
+import homeassistant.helpers.device_registry as dr
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
     async_dispatcher_send,
@@ -164,6 +165,8 @@ async def async_setup_entry(
         sw_version=gateway_info.firmware_version,
     )
 
+    remove_stale_devices(hass, entry, devices)
+
     # Setup the device coordinators
     coordinator_data = {
         CONF_GATEWAY_ID: gateway,
@@ -231,3 +234,45 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             listener()
 
     return unload_ok
+
+
+@callback
+def remove_stale_devices(
+    hass: HomeAssistant, config_entry: ConfigEntry, devices: list[Device]
+) -> None:
+    """Remove stale devices from device registry."""
+    device_registry = dr.async_get(hass)
+    device_entries = dr.async_entries_for_config_entry(
+        device_registry, config_entry.entry_id
+    )
+    all_device_ids = {device.id for device in devices}
+
+    for device_entry in device_entries:
+        device_id: str | None = None
+        gateway_id: str | None = None
+
+        for identifier in device_entry.identifiers:
+            if identifier[0] != DOMAIN:
+                continue
+
+            _id = identifier[1]
+
+            # Identify gateway device.
+            if _id == config_entry.data[CONF_GATEWAY_ID]:
+                gateway_id = _id
+                break
+
+            device_id = _id
+            break
+
+        if gateway_id is not None:
+            # Do not remove gateway device entry.
+            continue
+
+        if device_id is None or device_id not in all_device_ids:
+            # If device_id is None an invalid device entry was found for this config entry.
+            # If the device_id is not in existing device ids it's a stale device entry.
+            # Remove config entry from this device entry in either case.
+            device_registry.async_update_device(
+                device_entry.id, remove_config_entry_id=config_entry.entry_id
+            )
diff --git a/tests/components/tradfri/test_init.py b/tests/components/tradfri/test_init.py
index f96d1c09050..2a26391c43f 100644
--- a/tests/components/tradfri/test_init.py
+++ b/tests/components/tradfri/test_init.py
@@ -49,3 +49,44 @@ async def test_entry_setup_unload(hass, mock_api_factory):
         await hass.async_block_till_done()
         assert unload.call_count == len(tradfri.PLATFORMS)
         assert mock_api_factory.shutdown.call_count == 1
+
+
+async def test_remove_stale_devices(hass, mock_api_factory):
+    """Test remove stale device registry entries."""
+    entry = MockConfigEntry(
+        domain=tradfri.DOMAIN,
+        data={
+            tradfri.CONF_HOST: "mock-host",
+            tradfri.CONF_IDENTITY: "mock-identity",
+            tradfri.CONF_KEY: "mock-key",
+            tradfri.CONF_IMPORT_GROUPS: True,
+            tradfri.CONF_GATEWAY_ID: GATEWAY_ID,
+        },
+    )
+
+    entry.add_to_hass(hass)
+    dev_reg = dr.async_get(hass)
+    dev_reg.async_get_or_create(
+        config_entry_id=entry.entry_id,
+        identifiers={(tradfri.DOMAIN, "stale_device_id")},
+    )
+    dev_entries = dr.async_entries_for_config_entry(dev_reg, entry.entry_id)
+
+    assert len(dev_entries) == 1
+    dev_entry = dev_entries[0]
+    assert dev_entry.identifiers == {(tradfri.DOMAIN, "stale_device_id")}
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    dev_entries = dr.async_entries_for_config_entry(dev_reg, entry.entry_id)
+
+    # Check that only the gateway device entry remains.
+    assert len(dev_entries) == 1
+    dev_entry = dev_entries[0]
+    assert dev_entry.identifiers == {
+        (tradfri.DOMAIN, entry.data[tradfri.CONF_GATEWAY_ID])
+    }
+    assert dev_entry.manufacturer == tradfri.ATTR_TRADFRI_MANUFACTURER
+    assert dev_entry.name == tradfri.ATTR_TRADFRI_GATEWAY
+    assert dev_entry.model == tradfri.ATTR_TRADFRI_GATEWAY_MODEL