Simplify unifi cleanup logic (#65345)

homeassistant/components/unifi/unifi_entity_base.py

@@ -3,7 +3,7 @@ import logging
 from typing import Any
 
 from homeassistant.core import callback
-from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import Entity
 
@@ -78,34 +78,14 @@ class UniFiBase(Entity):
         raise NotImplementedError
 
     async def remove_item(self, keys: set) -> None:
-        """Remove entity if key is part of set.
+        """Remove entity if key is part of set."""
-
-        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:
             return
 
-        entity_registry = er.async_get(self.hass)
+        if self.registry_entry:
-        entity_entry = entity_registry.async_get(self.entity_id)
+            er.async_get(self.hass).async_remove(self.entity_id)
-        if not entity_entry:
+        else:
             await self.async_remove(force_remove=True)
-            return
-
-        device_registry = dr.async_get(self.hass)
-        device_entry = device_registry.async_get(entity_entry.device_id)
-        if not device_entry:
-            entity_registry.async_remove(self.entity_id)
-            return
-
-        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
     def should_poll(self) -> bool:

tests/components/unifi/test_device_tracker.py

@@ -18,7 +18,7 @@ from homeassistant.components.unifi.const import (
     DOMAIN as UNIFI_DOMAIN,
 )
 from homeassistant.const import STATE_HOME, STATE_NOT_HOME, STATE_UNAVAILABLE
-from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers import entity_registry as er
 import homeassistant.util.dt as dt_util
 
 from .test_controller import ENTRY_CONFIG, setup_unifi_integration
@@ -317,49 +317,6 @@ async def test_remove_clients(
     assert hass.states.get("device_tracker.client_2")
 
 
-async def test_remove_client_but_keep_device_entry(
-    hass, aioclient_mock, mock_unifi_websocket, mock_device_registry
-):
-    """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) == 3
-
-    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) == 2
-
-
 async def test_controller_state_change(
     hass, aioclient_mock, mock_unifi_websocket, mock_device_registry
 ):