From 57bcb4743dc2abbfbb90d36f3d948987c248a69e Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 19 May 2020 23:57:41 +0200 Subject: [PATCH] UniFi - Fix disabled entities being enabled after a restart (#35819) * Async remove call removed too much, resulting in disabled entities coming back after a restart * Calling super().async_remove is no longer needed, changed to self.async_remove * Yes, they should be sets... --- .../components/unifi/device_tracker.py | 8 ++-- homeassistant/components/unifi/sensor.py | 2 +- homeassistant/components/unifi/switch.py | 4 +- .../components/unifi/unifi_entity_base.py | 38 +++++++++---------- 4 files changed, 25 insertions(+), 27 deletions(-) diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 15c0ee7d4a7..ebad63acb4e 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -244,17 +244,17 @@ class UniFiClientTracker(UniFiClient, ScannerEntity): async def options_updated(self) -> None: """Config entry options are updated, remove entity if option is disabled.""" if not self.controller.option_track_clients: - await self.async_remove() + await self.remove_item({self.client.mac}) elif self.is_wired: if not self.controller.option_track_wired_clients: - await self.async_remove() + await self.remove_item({self.client.mac}) elif ( self.controller.option_ssid_filter and self.client.essid not in self.controller.option_ssid_filter ): - await self.async_remove() + await self.remove_item({self.client.mac}) class UniFiDeviceTracker(UniFiBase, ScannerEntity): @@ -383,4 +383,4 @@ class UniFiDeviceTracker(UniFiBase, ScannerEntity): async def options_updated(self) -> None: """Config entry options are updated, remove entity if option is disabled.""" if not self.controller.option_track_devices: - await self.async_remove() + await self.remove_item({self.device.mac}) diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index cf3cc2b1b5a..8fdb0ac1461 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -73,7 +73,7 @@ class UniFiBandwidthSensor(UniFiClient): async def options_updated(self) -> None: """Config entry options are updated, remove entity if option is disabled.""" if not self.controller.option_allow_bandwidth_sensors: - await self.async_remove() + await self.remove_item({self.client.mac}) class UniFiRxBandwidthSensor(UniFiBandwidthSensor): diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index a216daaf9bb..1b9f7147a9a 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -239,7 +239,7 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity): async def options_updated(self) -> None: """Config entry options are updated, remove entity if option is disabled.""" if not self.controller.option_poe_clients: - await self.async_remove() + await self.remove_item({self.client.mac}) class UniFiBlockClientSwitch(UniFiClient, SwitchEntity): @@ -287,4 +287,4 @@ class UniFiBlockClientSwitch(UniFiClient, SwitchEntity): async def options_updated(self) -> None: """Config entry options are updated, remove entity if option is disabled.""" if self.client.mac not in self.controller.option_block_clients: - await self.async_remove() + await self.remove_item({self.client.mac}) diff --git a/homeassistant/components/unifi/unifi_entity_base.py b/homeassistant/components/unifi/unifi_entity_base.py index 6fbc18f75d3..46a7123e4c9 100644 --- a/homeassistant/components/unifi/unifi_entity_base.py +++ b/homeassistant/components/unifi/unifi_entity_base.py @@ -43,18 +43,33 @@ class UniFiBase(Entity): self._item.remove_callback(self.async_update_callback) self.controller.entities[self.DOMAIN][self.TYPE].remove(self._item.mac) - async def async_remove(self): - """Clean up when removing entity. + @callback + def async_update_callback(self) -> None: + """Update the entity's state.""" + LOGGER.debug( + "Updating %s entity %s (%s)", self.TYPE, self.entity_id, self._item.mac + ) + self.async_write_ha_state() + + 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. 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. """ + if self._item.mac not in mac_addresses: + return + 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() + await self.async_remove() return device_registry = await self.hass.helpers.device_registry.async_get_registry() @@ -69,23 +84,6 @@ class UniFiBase(Entity): entity_registry.async_remove(self.entity_id) - @callback - def async_update_callback(self) -> None: - """Update the entity's state.""" - LOGGER.debug( - "Updating %s entity %s (%s)", self.TYPE, self.entity_id, self._item.mac - ) - self.async_write_ha_state() - - 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._item.mac in mac_addresses: - await self.async_remove() - @property def should_poll(self) -> bool: """No polling needed."""