From ca6b53c16dd49d8f6ebacfd29cb94e8a2706bf1c Mon Sep 17 00:00:00 2001
From: Robert Svensson <Kane610@users.noreply.github.com>
Date: Mon, 27 Sep 2021 12:04:29 +0200
Subject: [PATCH] Remove UniFi config entry reference from device when removing
 last entity of said device (#56501)

---
 .../components/unifi/unifi_entity_base.py     | 23 ++++++++--
 tests/components/unifi/test_device_tracker.py | 43 +++++++++++++++++++
 2 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/unifi/unifi_entity_base.py b/homeassistant/components/unifi/unifi_entity_base.py
index 03c63ce4e84..9d2d8071fca 100644
--- a/homeassistant/components/unifi/unifi_entity_base.py
+++ b/homeassistant/components/unifi/unifi_entity_base.py
@@ -83,6 +83,7 @@ class UniFiBase(Entity):
         Remove entity if no entry in entity registry exist.
         Remove entity registry entry if no entry in device registry exist.
         Remove device registry entry if there is only one linked entity (this entity).
+        Remove config entry reference from device registry entry if there is more than one config entry.
         Remove entity registry entry if there are more than one entity linked to the device registry entry.
         """
         if self.key not in keys:
@@ -102,17 +103,33 @@ class UniFiBase(Entity):
 
         if (
             len(
-                async_entries_for_device(
+                entries_for_device := async_entries_for_device(
                     entity_registry,
                     entity_entry.device_id,
                     include_disabled_entities=True,
                 )
             )
-            == 1
-        ):
+        ) == 1:
             device_registry.async_remove_device(device_entry.id)
             return
 
+        if (
+            len(
+                entries_for_device_from_this_config_entry := [
+                    entry_for_device
+                    for entry_for_device in entries_for_device
+                    if entry_for_device.config_entry_id
+                    == self.controller.config_entry.entry_id
+                ]
+            )
+            != len(entries_for_device)
+            and len(entries_for_device_from_this_config_entry) == 1
+        ):
+            device_registry.async_update_device(
+                entity_entry.device_id,
+                remove_config_entry_id=self.controller.config_entry.entry_id,
+            )
+
         entity_registry.async_remove(self.entity_id)
 
     @property
diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py
index d583cad86c3..4ba690dc444 100644
--- a/tests/components/unifi/test_device_tracker.py
+++ b/tests/components/unifi/test_device_tracker.py
@@ -381,6 +381,49 @@ async def test_remove_clients(hass, aioclient_mock, mock_unifi_websocket):
     assert hass.states.get("device_tracker.client_2")
 
 
+async def test_remove_client_but_keep_device_entry(
+    hass, aioclient_mock, mock_unifi_websocket
+):
+    """Test that unifi entity base remove config entry id from a multi integration device registry entry."""
+    client_1 = {
+        "essid": "ssid",
+        "hostname": "client_1",
+        "is_wired": False,
+        "last_seen": 1562600145,
+        "mac": "00:00:00:00:00:01",
+    }
+    await setup_unifi_integration(hass, aioclient_mock, clients_response=[client_1])
+
+    device_registry = dr.async_get(hass)
+    device_entry = device_registry.async_get_or_create(
+        config_entry_id="other",
+        connections={("mac", "00:00:00:00:00:01")},
+    )
+
+    entity_registry = er.async_get(hass)
+    other_entity = entity_registry.async_get_or_create(
+        TRACKER_DOMAIN,
+        "other",
+        "unique_id",
+        device_id=device_entry.id,
+    )
+    assert len(device_entry.config_entries) == 2
+
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_CLIENT_REMOVED},
+            "data": [client_1],
+        }
+    )
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 0
+
+    device_entry = device_registry.async_get(other_entity.device_id)
+    assert len(device_entry.config_entries) == 1
+
+
 async def test_controller_state_change(hass, aioclient_mock, mock_unifi_websocket):
     """Verify entities state reflect on controller becoming unavailable."""
     client = {