UniFi - Change handling of updated options (#31762)

* Change handling of updated options

* Add tests

homeassistant/components/unifi/controller.py

@@ -164,7 +164,7 @@ class UniFiController:
                     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 @@ class UniFiController:
 
         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):
+    async def async_config_entry_updated(hass, entry) -> None:
-        """Triggered by config entry options updates."""
+        """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.config_entry import ScannerEntity
 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,6 +67,12 @@ 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(
@@ -72,24 +80,59 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     )
 
     @callback
-    def update_disable_on_entities():
+    def options_updated():
-        """Update the values of the controller."""
+        """Manage entities affected by config entry options."""
-        for entity in tracked.values():
+        nonlocal option_track_clients
+        nonlocal option_track_devices
+        nonlocal option_track_wired_clients
 
-            if entity.entity_registry_enabled_default == entity.enabled:
+        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 config_entry_option:
-            if not entity.entity_registry_enabled_default and entity.enabled:
+                update = True
-                disabled_by = DISABLED_CONFIG_ENTRY
+            else:
+                for mac, entity in tracked.items():
+                    if isinstance(entity, tracker_class):
+                        remove.add(mac)
 
-            registry.async_update_entity(
+        if (
-                entity.registry_entry.entity_id, disabled_by=disabled_by
+            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 (
+    for items, tracker_class, track in (
-        (controller.api.clients, UniFiClientTracker),
+        (controller.api.clients, UniFiClientTracker, controller.option_track_clients),
-        (controller.api.devices, UniFiDeviceTracker),
+        (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,11 +180,12 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
             self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time
 
     @property
-    def entity_registry_enabled_default(self):
+    def is_connected(self):
-        """Return if the entity should be enabled when first added to the entity registry."""
+        """Return true if the client is connected to the network.
-        if not self.controller.option_track_clients:
-            return False
 
+        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
@@ -142,17 +193,6 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         ):
             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 @@ class UniFiDeviceTracker(ScannerEntity):
         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 @@ import logging
 
 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,11 +22,18 @@ 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(
@@ -36,24 +41,29 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     )
 
     @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:
+        if option_allow_bandwidth_sensors != controller.option_allow_bandwidth_sensors:
-                continue
+            option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors
 
-            disabled_by = None
+            if option_allow_bandwidth_sensors:
-            if not entity.entity_registry_enabled_default and entity.enabled:
+                update_controller()
-                disabled_by = DISABLED_CONFIG_ENTRY
 
-            registry.async_update_entity(
+            else:
-                entity.registry_entry.entity_id, disabled_by=disabled_by
+                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."""

tests/components/unifi/test_device_tracker.py

@@ -11,6 +11,7 @@ from homeassistant.components import unifi
 import homeassistant.components.device_tracker as device_tracker
 from homeassistant.components.unifi.const import (
     CONF_SSID_FILTER,
+    CONF_TRACK_CLIENTS,
     CONF_TRACK_DEVICES,
     CONF_TRACK_WIRED_CLIENTS,
 )
@@ -114,7 +115,7 @@ async def test_tracked_devices(hass):
         devices_response=[DEVICE_1, DEVICE_2],
         known_wireless_clients=(CLIENT_4["mac"],),
     )
-    assert len(hass.states.async_all()) == 5
+    assert len(hass.states.async_all()) == 6
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
@@ -125,7 +126,8 @@ async def test_tracked_devices(hass):
     assert client_2.state == "not_home"
 
     client_3 = hass.states.get("device_tracker.client_3")
-    assert client_3 is None
+    assert client_3 is not None
+    assert client_3.state == "not_home"
 
     # Wireless client with wired bug, if bug active on restart mark device away
     client_4 = hass.states.get("device_tracker.client_4")
@@ -136,6 +138,7 @@ async def test_tracked_devices(hass):
     assert device_1 is not None
     assert device_1.state == "not_home"
 
+    # State change signalling works
     client_1_copy = copy(CLIENT_1)
     client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": "sta:sync"}, "data": [client_1_copy]}
@@ -152,28 +155,6 @@ async def test_tracked_devices(hass):
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1.state == "home"
 
-    # Controller unavailable
-    controller.async_unifi_signalling_callback(
-        SIGNAL_CONNECTION_STATE, STATE_DISCONNECTED
-    )
-    await hass.async_block_till_done()
-
-    client_1 = hass.states.get("device_tracker.client_1")
-    assert client_1.state == STATE_UNAVAILABLE
-
-    device_1 = hass.states.get("device_tracker.device_1")
-    assert device_1.state == STATE_UNAVAILABLE
-
-    # Controller available
-    controller.async_unifi_signalling_callback(SIGNAL_CONNECTION_STATE, STATE_RUNNING)
-    await hass.async_block_till_done()
-
-    client_1 = hass.states.get("device_tracker.client_1")
-    assert client_1.state == "home"
-
-    device_1 = hass.states.get("device_tracker.device_1")
-    assert device_1.state == "home"
-
     # Disabled device is unavailable
     device_1_copy = copy(DEVICE_1)
     device_1_copy["disabled"] = True
@@ -184,24 +165,205 @@ async def test_tracked_devices(hass):
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1.state == STATE_UNAVAILABLE
 
-    # Don't track wired clients nor devices
+
-    controller.config_entry.add_update_listener(controller.async_options_updated)
+async def test_controller_state_change(hass):
-    hass.config_entries.async_update_entry(
+    """Verify entities state reflect on controller becoming unavailable."""
-        controller.config_entry,
+    controller = await setup_unifi_integration(
-        options={
+        hass, clients_response=[CLIENT_1], devices_response=[DEVICE_1],
-            CONF_SSID_FILTER: [],
+    )
-            CONF_TRACK_WIRED_CLIENTS: False,
+    assert len(hass.states.async_all()) == 3
-            CONF_TRACK_DEVICES: False,
+
-        },
+    # Controller unavailable
+    controller.async_unifi_signalling_callback(
+        SIGNAL_CONNECTION_STATE, STATE_DISCONNECTED
     )
     await hass.async_block_till_done()
+
     client_1 = hass.states.get("device_tracker.client_1")
-    assert client_1
+    assert client_1.state == STATE_UNAVAILABLE
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1.state == STATE_UNAVAILABLE
+
+    # Controller available
+    controller.async_unifi_signalling_callback(SIGNAL_CONNECTION_STATE, STATE_RUNNING)
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "not_home"
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1.state == "not_home"
+
+
+async def test_option_track_clients(hass):
+    """Test the tracking of clients can be turned off."""
+    controller = await setup_unifi_integration(
+        hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1],
+    )
+    assert len(hass.states.async_all()) == 4
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_CLIENTS: False},
+    )
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is None
+
     client_2 = hass.states.get("device_tracker.wired_client")
     assert client_2 is None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_CLIENTS: True},
+    )
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+
+async def test_option_track_wired_clients(hass):
+    """Test the tracking of wired clients can be turned off."""
+    controller = await setup_unifi_integration(
+        hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1],
+    )
+    assert len(hass.states.async_all()) == 4
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_WIRED_CLIENTS: False},
+    )
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_WIRED_CLIENTS: True},
+    )
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+
+async def test_option_track_devices(hass):
+    """Test the tracking of devices can be turned off."""
+    controller = await setup_unifi_integration(
+        hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1],
+    )
+    assert len(hass.states.async_all()) == 4
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_DEVICES: False},
+    )
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None
+
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1 is None
 
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_TRACK_DEVICES: True},
+    )
+    await hass.async_block_till_done()
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1 is not None
+
+    client_2 = hass.states.get("device_tracker.wired_client")
+    assert client_2 is not None
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1 is not None
+
+
+async def test_option_ssid_filter(hass):
+    """Test the SSID filter works."""
+    controller = await setup_unifi_integration(
+        hass, options={CONF_SSID_FILTER: ["ssid"]}, clients_response=[CLIENT_3],
+    )
+    assert len(hass.states.async_all()) == 2
+
+    # SSID filter active
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3.state == "not_home"
+
+    client_3_copy = copy(CLIENT_3)
+    client_3_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    event = {"meta": {"message": "sta:sync"}, "data": [client_3_copy]}
+    controller.api.message_handler(event)
+    await hass.async_block_till_done()
+
+    # SSID filter active even though time stamp should mark as home
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3.state == "not_home"
+
+    # Remove SSID filter
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_SSID_FILTER: []},
+    )
+    event = {"meta": {"message": "sta:sync"}, "data": [client_3_copy]}
+    controller.api.message_handler(event)
+    await hass.async_block_till_done()
+
+    # SSID no longer filtered
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3.state == "home"
+
 
 async def test_wireless_client_go_wired_issue(hass):
     """Test the solution to catch wireless device go wired UniFi issue.

tests/components/unifi/test_sensor.py

@@ -99,3 +99,27 @@ async def test_sensors(hass):
 
     wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
     assert wireless_client_tx.state == "6789.0"
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry,
+        options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: False},
+    )
+    await hass.async_block_till_done()
+
+    wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
+    assert wireless_client_rx is None
+
+    wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
+    assert wireless_client_tx is None
+
+    hass.config_entries.async_update_entry(
+        controller.config_entry,
+        options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
+    )
+    await hass.async_block_till_done()
+
+    wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
+    assert wireless_client_rx.state == "2345.0"
+
+    wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
+    assert wireless_client_tx.state == "6789.0"