diff --git a/homeassistant/components/unifi/unifi_entity_base.py b/homeassistant/components/unifi/unifi_entity_base.py index 1c3251d213c..c611fc1ee60 100644 --- a/homeassistant/components/unifi/unifi_entity_base.py +++ b/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 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. - """ + """Remove entity if key is part of set.""" if self.key not in keys: return - entity_registry = er.async_get(self.hass) - entity_entry = entity_registry.async_get(self.entity_id) - if not entity_entry: + if self.registry_entry: + er.async_get(self.hass).async_remove(self.entity_id) + 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: diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index cf4861980a0..b490d43fffd 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/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 ):