UniFi - refactor entity management (#34367)

* Move removal of sensor entities into a base class

* Fix martins comments on sensors

* Reflect sensor changes on device_tracker platform

* Reflect sensor changes on switch platform

* Improve layering

* Make sure to clean up entity and device registry when removing entities

* Fix martins comments

homeassistant/components/unifi/unifi_entity_base.py

@@ -0,0 +1,80 @@
+"""Base class for UniFi entities."""
+
+from homeassistant.core import callback
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.entity_registry import async_entries_for_device
+
+
+class UniFiBase(Entity):
+    """UniFi entity base class."""
+
+    TYPE = ""
+
+    def __init__(self, controller) -> None:
+        """Set up UniFi entity base."""
+        self.controller = controller
+
+    @property
+    def mac(self):
+        """Return MAC of entity."""
+        raise NotImplementedError
+
+    async def async_added_to_hass(self) -> None:
+        """Entity created."""
+        self.controller.entities[self.platform.domain][self.TYPE].add(self.mac)
+        for signal, method in (
+            (self.controller.signal_reachable, self.async_update_callback),
+            (self.controller.signal_options_update, self.options_updated),
+            (self.controller.signal_remove, self.remove_item),
+        ):
+            self.async_on_remove(async_dispatcher_connect(self.hass, signal, method))
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Disconnect object when removed."""
+        self.controller.entities[self.platform.domain][self.TYPE].remove(self.mac)
+
+    async def async_remove(self):
+        """Clean up when removing 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 entity registry entry if there are more than one entity linked to the device registry entry.
+        """
+        entity_registry = await self.hass.helpers.entity_registry.async_get_registry()
+        entity_entry = entity_registry.async_get(self.entity_id)
+        if not entity_entry:
+            await super().async_remove()
+            return
+
+        device_registry = await self.hass.helpers.device_registry.async_get_registry()
+        device_entry = device_registry.async_get(entity_entry.device_id)
+        if not device_entry:
+            entity_registry.async_remove(self.entity_id)
+            return
+
+        if len(async_entries_for_device(entity_registry, entity_entry.device_id)) == 1:
+            device_registry.async_remove_device(device_entry.id)
+            return
+
+        entity_registry.async_remove(self.entity_id)
+
+    @callback
+    def async_update_callback(self):
+        """Update the entity's state."""
+        raise NotImplementedError
+
+    async def options_updated(self) -> None:
+        """Config entry options are updated, remove entity if option is disabled."""
+        raise NotImplementedError
+
+    async def remove_item(self, mac_addresses: set) -> None:
+        """Remove entity if MAC is part of set."""
+        if self.mac in mac_addresses:
+            await self.async_remove()
+
+    @property
+    def should_poll(self) -> bool:
+        """No polling needed."""
+        return True