UniFi - Change handling of updated options (#31762)

* Change handling of updated options

* Add tests

homeassistant/components/unifi/controller.py

@@ -164,7 +164,7 @@ def async_unifi_signalling_callback(self, signal, data):
                     WIRELESS_GUEST_CONNECTED,
                 ):
                     self.update_wireless_clients()
-            else:
+            elif data.get("clients") or data.get("devices"):
                 async_dispatcher_send(self.hass, self.signal_update)
 
     @property
@@ -238,13 +238,13 @@ async def async_setup(self):
 
         self.api.start_websocket()
 
-        self.config_entry.add_update_listener(self.async_options_updated)
+        self.config_entry.add_update_listener(self.async_config_entry_updated)
 
         return True
 
     @staticmethod
-    async def async_options_updated(hass, entry):
-        """Triggered by config entry options updates."""
+    async def async_config_entry_updated(hass, entry) -> None:
+        """Handle signals of config entry being updated."""
         controller_id = CONTROLLER_ID.format(
             host=entry.data[CONF_CONTROLLER][CONF_HOST],
             site=entry.data[CONF_CONTROLLER][CONF_SITE_ID],

homeassistant/components/unifi/device_tracker.py

@@ -6,10 +6,8 @@
 from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
 from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
 from homeassistant.core import callback
-from homeassistant.helpers import entity_registry
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY
 import homeassistant.util.dt as dt_util
 
 from .const import ATTR_MANUFACTURER
@@ -43,7 +41,11 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     controller = get_controller_from_config_entry(hass, config_entry)
     tracked = {}
 
-    registry = await entity_registry.async_get_registry(hass)
+    option_track_clients = controller.option_track_clients
+    option_track_devices = controller.option_track_devices
+    option_track_wired_clients = controller.option_track_wired_clients
+
+    registry = await hass.helpers.entity_registry.async_get_registry()
 
     # Restore clients that is not a part of active clients list.
     for entity in registry.entities.values():
@@ -65,31 +67,72 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     @callback
     def update_controller():
         """Update the values of the controller."""
+        nonlocal option_track_clients
+        nonlocal option_track_devices
+
+        if not option_track_clients and not option_track_devices:
+            return
+
         add_entities(controller, async_add_entities, tracked)
 
     controller.listeners.append(
         async_dispatcher_connect(hass, controller.signal_update, update_controller)
     )
 
     @callback
-    def update_disable_on_entities():
-        """Update the values of the controller."""
-        for entity in tracked.values():
-
-            if entity.entity_registry_enabled_default == entity.enabled:
+    def options_updated():
+        """Manage entities affected by config entry options."""
+        nonlocal option_track_clients
+        nonlocal option_track_devices
+        nonlocal option_track_wired_clients
+
+        update = False
+        remove = set()
+
+        for current_option, config_entry_option, tracker_class in (
+            (option_track_clients, controller.option_track_clients, UniFiClientTracker),
+            (option_track_devices, controller.option_track_devices, UniFiDeviceTracker),
+        ):
+            if current_option == config_entry_option:
                 continue
 
-            disabled_by = None
-            if not entity.entity_registry_enabled_default and entity.enabled:
-                disabled_by = DISABLED_CONFIG_ENTRY
+            if config_entry_option:
+                update = True
+            else:
+                for mac, entity in tracked.items():
+                    if isinstance(entity, tracker_class):
+                        remove.add(mac)
 
-            registry.async_update_entity(
-                entity.registry_entry.entity_id, disabled_by=disabled_by
-            )
+        if (
+            controller.option_track_clients
+            and option_track_wired_clients != controller.option_track_wired_clients
+        ):
+
+            if controller.option_track_wired_clients:
+                update = True
+            else:
+                for mac, entity in tracked.items():
+                    if isinstance(entity, UniFiClientTracker) and entity.is_wired:
+                        remove.add(mac)
+
+        option_track_clients = controller.option_track_clients
+        option_track_devices = controller.option_track_devices
+        option_track_wired_clients = controller.option_track_wired_clients
+
+        for mac in remove:
+            entity = tracked.pop(mac)
+
+            if registry.async_is_registered(entity.entity_id):
+                registry.async_remove(entity.entity_id)
+
+            hass.async_create_task(entity.async_remove())
+
+        if update:
+            update_controller()
 
     controller.listeners.append(
         async_dispatcher_connect(
-            hass, controller.signal_options_update, update_disable_on_entities
+            hass, controller.signal_options_update, options_updated
         )
     )
 
@@ -101,16 +144,23 @@ def add_entities(controller, async_add_entities, tracked):
     """Add new tracker entities from the controller."""
     new_tracked = []
 
-    for items, tracker_class in (
-        (controller.api.clients, UniFiClientTracker),
-        (controller.api.devices, UniFiDeviceTracker),
+    for items, tracker_class, track in (
+        (controller.api.clients, UniFiClientTracker, controller.option_track_clients),
+        (controller.api.devices, UniFiDeviceTracker, controller.option_track_devices),
     ):
+        if not track:
+            continue
 
         for item_id in items:
 
             if item_id in tracked:
                 continue
 
+            if tracker_class is UniFiClientTracker and (
+                not controller.option_track_wired_clients and items[item_id].is_wired
+            ):
+                continue
+
             tracked[item_id] = tracker_class(items[item_id], controller)
             new_tracked.append(tracked[item_id])
 
@@ -130,29 +180,19 @@ def __init__(self, client, controller):
             self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time
 
     @property
-    def entity_registry_enabled_default(self):
-        """Return if the entity should be enabled when first added to the entity registry."""
-        if not self.controller.option_track_clients:
-            return False
+    def is_connected(self):
+        """Return true if the client is connected to the network.
 
+        If connected to unwanted ssid return False.
+        If is_wired and client.is_wired differ it means that the device is offline and UniFi bug shows device as wired.
+        """
         if (
             not self.is_wired
             and self.controller.option_ssid_filter
             and self.client.essid not in self.controller.option_ssid_filter
         ):
             return False
 
-        if not self.controller.option_track_wired_clients and self.is_wired:
-            return False
-
-        return True
-
-    @property
-    def is_connected(self):
-        """Return true if the client is connected to the network.
-
-        If is_wired and client.is_wired differ it means that the device is offline and UniFi bug shows device as wired.
-        """
         if self.is_wired != self.client.is_wired:
             if not self.wired_bug:
                 self.wired_bug = dt_util.utcnow()
@@ -202,13 +242,6 @@ def __init__(self, device, controller):
         self.controller = controller
         self.listeners = []
 
-    @property
-    def entity_registry_enabled_default(self):
-        """Return if the entity should be enabled when first added to the entity registry."""
-        if self.controller.option_track_devices:
-            return True
-        return False
-
     async def async_added_to_hass(self):
         """Subscribe to device events."""
         LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac)

homeassistant/components/unifi/sensor.py

@@ -3,9 +3,7 @@
 
 from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
 from homeassistant.core import callback
-from homeassistant.helpers import entity_registry
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY
 
 from .unifi_client import UniFiClient
 
@@ -24,36 +22,48 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     controller = get_controller_from_config_entry(hass, config_entry)
     sensors = {}
 
-    registry = await entity_registry.async_get_registry(hass)
+    option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors
+
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
 
     @callback
     def update_controller():
         """Update the values of the controller."""
+        nonlocal option_allow_bandwidth_sensors
+
+        if not option_allow_bandwidth_sensors:
+            return
+
         add_entities(controller, async_add_entities, sensors)
 
     controller.listeners.append(
         async_dispatcher_connect(hass, controller.signal_update, update_controller)
     )
 
     @callback
-    def update_disable_on_entities():
+    def options_updated():
         """Update the values of the controller."""
-        for entity in sensors.values():
+        nonlocal option_allow_bandwidth_sensors
 
-            if entity.entity_registry_enabled_default == entity.enabled:
-                continue
+        if option_allow_bandwidth_sensors != controller.option_allow_bandwidth_sensors:
+            option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors
 
-            disabled_by = None
-            if not entity.entity_registry_enabled_default and entity.enabled:
-                disabled_by = DISABLED_CONFIG_ENTRY
+            if option_allow_bandwidth_sensors:
+                update_controller()
 
-            registry.async_update_entity(
-                entity.registry_entry.entity_id, disabled_by=disabled_by
-            )
+            else:
+                for sensor in sensors.values():
+
+                    if entity_registry.async_is_registered(sensor.entity_id):
+                        entity_registry.async_remove(sensor.entity_id)
+
+                    hass.async_create_task(sensor.async_remove())
+
+                sensors.clear()
 
     controller.listeners.append(
         async_dispatcher_connect(
-            hass, controller.signal_options_update, update_disable_on_entities
+            hass, controller.signal_options_update, options_updated
         )
     )
 
@@ -87,13 +97,6 @@ def add_entities(controller, async_add_entities, sensors):
 class UniFiRxBandwidthSensor(UniFiClient):
     """Receiving bandwidth sensor."""
 
-    @property
-    def entity_registry_enabled_default(self):
-        """Return if the entity should be enabled when first added to the entity registry."""
-        if self.controller.option_allow_bandwidth_sensors:
-            return True
-        return False
-
     @property
     def state(self):
         """Return the state of the sensor."""