UniFi - use entity registry disabled_by to control available entities (#26141)

* Move ignoring logic to entity registry enabled default

* Handle config to option import better

* Properly enable and disable entity registry entries on changes from config entry options

* Fix balloobs comments

* Fix some tests

* Fix tests

* Simplify updating disable on entities

* Simplify device tracker update function

* Local entity disabled replaced on rebase

* Only alter entities with changed options

* Proper tracking of changed options

* Back to straightforward updating of disabled

homeassistant/components/unifi/__init__.py

@@ -11,6 +11,9 @@ from .const import (
     CONF_BLOCK_CLIENT,
     CONF_CONTROLLER,
     CONF_DETECTION_TIME,
+    CONF_DONT_TRACK_CLIENTS,
+    CONF_DONT_TRACK_DEVICES,
+    CONF_DONT_TRACK_WIRED_CLIENTS,
     CONF_SITE_ID,
     CONF_SSID_FILTER,
     CONTROLLER_ID,
@@ -20,9 +23,6 @@ from .const import (
 from .controller import UniFiController
 
 CONF_CONTROLLERS = "controllers"
-CONF_DONT_TRACK_CLIENTS = "dont_track_clients"
-CONF_DONT_TRACK_DEVICES = "dont_track_devices"
-CONF_DONT_TRACK_WIRED_CLIENTS = "dont_track_wired_clients"
 
 CONTROLLER_SCHEMA = vol.Schema(
     {

homeassistant/components/unifi/const.py

@@ -18,6 +18,10 @@ CONF_TRACK_DEVICES = "track_devices"
 CONF_TRACK_WIRED_CLIENTS = "track_wired_clients"
 CONF_SSID_FILTER = "ssid_filter"
 
+CONF_DONT_TRACK_CLIENTS = "dont_track_clients"
+CONF_DONT_TRACK_DEVICES = "dont_track_devices"
+CONF_DONT_TRACK_WIRED_CLIENTS = "dont_track_wired_clients"
+
 DEFAULT_BLOCK_CLIENTS = []
 DEFAULT_TRACK_CLIENTS = True
 DEFAULT_TRACK_DEVICES = True

homeassistant/components/unifi/controller.py

@@ -18,6 +18,9 @@ from .const import (
     CONF_BLOCK_CLIENT,
     CONF_CONTROLLER,
     CONF_DETECTION_TIME,
+    CONF_DONT_TRACK_CLIENTS,
+    CONF_DONT_TRACK_DEVICES,
+    CONF_DONT_TRACK_WIRED_CLIENTS,
     CONF_TRACK_CLIENTS,
     CONF_TRACK_DEVICES,
     CONF_TRACK_WIRED_CLIENTS,
@@ -30,6 +33,7 @@ from .const import (
     DEFAULT_TRACK_WIRED_CLIENTS,
     DEFAULT_DETECTION_TIME,
     DEFAULT_SSID_FILTER,
+    DOMAIN,
     LOGGER,
     UNIFI_CONFIG,
 )
@@ -49,7 +53,6 @@ class UniFiController:
 
         self._site_name = None
         self._site_role = None
-        self.unifi_config = {}
 
     @property
     def host(self):
@@ -116,11 +119,14 @@ class UniFiController:
         return None
 
     @property
-    def event_update(self):
+    def signal_update(self):
         """Event specific per UniFi entry to signal new data."""
-        return "unifi-update-{}".format(
+        return f"unifi-update-{CONTROLLER_ID.format(host=self.host, site=self.site)}"
-            CONTROLLER_ID.format(host=self.host, site=self.site)
+
-        )
+    @property
+    def signal_options_update(self):
+        """Event specific per UniFi entry to signal new options."""
+        return f"unifi-options-{CONTROLLER_ID.format(host=self.host, site=self.site)}"
 
     async def request_update(self):
         """Request an update."""
@@ -164,7 +170,7 @@ class UniFiController:
             LOGGER.info("Reconnected to controller %s", self.host)
             self.available = True
 
-        async_dispatcher_send(self.hass, self.event_update)
+        async_dispatcher_send(self.hass, self.signal_update)
 
     async def async_setup(self):
         """Set up a UniFi controller."""
@@ -191,37 +197,9 @@ class UniFiController:
             LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
             return False
 
-        for unifi_config in hass.data[UNIFI_CONFIG]:
+        self.import_configuration()
-            if (
-                self.host == unifi_config[CONF_HOST]
-                and self.site_name == unifi_config[CONF_SITE_ID]
-            ):
-                self.unifi_config = unifi_config
-                break
 
-        options = dict(self.config_entry.options)
+        self.config_entry.add_update_listener(self.async_options_updated)
-
-        if CONF_BLOCK_CLIENT in self.unifi_config:
-            options[CONF_BLOCK_CLIENT] = self.unifi_config[CONF_BLOCK_CLIENT]
-
-        if CONF_TRACK_CLIENTS in self.unifi_config:
-            options[CONF_TRACK_CLIENTS] = self.unifi_config[CONF_TRACK_CLIENTS]
-
-        if CONF_TRACK_DEVICES in self.unifi_config:
-            options[CONF_TRACK_DEVICES] = self.unifi_config[CONF_TRACK_DEVICES]
-
-        if CONF_TRACK_WIRED_CLIENTS in self.unifi_config:
-            options[CONF_TRACK_WIRED_CLIENTS] = self.unifi_config[
-                CONF_TRACK_WIRED_CLIENTS
-            ]
-
-        if CONF_DETECTION_TIME in self.unifi_config:
-            options[CONF_DETECTION_TIME] = self.unifi_config[CONF_DETECTION_TIME]
-
-        if CONF_SSID_FILTER in self.unifi_config:
-            options[CONF_SSID_FILTER] = self.unifi_config[CONF_SSID_FILTER]
-
-        hass.config_entries.async_update_entry(self.config_entry, options=options)
 
         for platform in ["device_tracker", "switch"]:
             hass.async_create_task(
@@ -232,6 +210,56 @@ class UniFiController:
 
         return True
 
+    @staticmethod
+    async def async_options_updated(hass, entry):
+        """Triggered by config entry options updates."""
+        controller_id = CONTROLLER_ID.format(
+            host=entry.data[CONF_CONTROLLER][CONF_HOST],
+            site=entry.data[CONF_CONTROLLER][CONF_SITE_ID],
+        )
+        controller = hass.data[DOMAIN][controller_id]
+
+        async_dispatcher_send(hass, controller.signal_options_update)
+
+    def import_configuration(self):
+        """Import configuration to config entry options."""
+        unifi_config = {}
+        for config in self.hass.data[UNIFI_CONFIG]:
+            if (
+                self.host == config[CONF_HOST]
+                and self.site_name == config[CONF_SITE_ID]
+            ):
+                unifi_config = config
+                break
+
+        old_options = dict(self.config_entry.options)
+        new_options = {}
+
+        for config, option in (
+            (CONF_BLOCK_CLIENT, CONF_BLOCK_CLIENT),
+            (CONF_DONT_TRACK_CLIENTS, CONF_TRACK_CLIENTS),
+            (CONF_DONT_TRACK_WIRED_CLIENTS, CONF_TRACK_WIRED_CLIENTS),
+            (CONF_DONT_TRACK_DEVICES, CONF_TRACK_DEVICES),
+            (CONF_DETECTION_TIME, CONF_DETECTION_TIME),
+            (CONF_SSID_FILTER, CONF_SSID_FILTER),
+        ):
+            if config in unifi_config:
+                if config == option and unifi_config[
+                    config
+                ] != self.config_entry.options.get(option):
+                    new_options[option] = unifi_config[config]
+                elif config != option and (
+                    option not in self.config_entry.options
+                    or unifi_config[config] == self.config_entry.options.get(option)
+                ):
+                    new_options[option] = not unifi_config[config]
+
+        if new_options:
+            options = {**old_options, **new_options}
+            self.hass.config_entries.async_update_entry(
+                self.config_entry, options=options
+            )
+
     async def async_reset(self):
         """Reset this controller to default state.
 

homeassistant/components/unifi/device_tracker.py

@@ -20,6 +20,7 @@ from homeassistant.const import (
 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.helpers.config_validation as cv
 import homeassistant.util.dt as dt_util
@@ -136,7 +137,24 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         """Update the values of the controller."""
         update_items(controller, async_add_entities, tracked)
 
-    async_dispatcher_connect(hass, controller.event_update, update_controller)
+    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():
+
+            disabled_by = None
+            if not entity.entity_registry_enabled_default and entity.enabled:
+                disabled_by = DISABLED_CONFIG_ENTRY
+
+            registry.async_update_entity(
+                entity.registry_entry.entity_id, disabled_by=disabled_by
+            )
+
+    async_dispatcher_connect(
+        hass, controller.signal_options_update, update_disable_on_entities
+    )
 
     update_controller()
 
@@ -146,65 +164,20 @@ def update_items(controller, async_add_entities, tracked):
     """Update tracked device state from the controller."""
     new_tracked = []
 
-    if controller.option_track_clients:
+    for items, tracker_class in (
+        (controller.api.clients, UniFiClientTracker),
+        (controller.api.devices, UniFiDeviceTracker),
+    ):
 
-        for client_id in controller.api.clients:
+        for item_id in items:
 
-            if client_id in tracked:
+            if item_id in tracked:
-                if not tracked[client_id].enabled:
+                if tracked[item_id].enabled:
-                    continue
+                    tracked[item_id].async_schedule_update_ha_state()
-                LOGGER.debug(
-                    "Updating UniFi tracked client %s (%s)",
-                    tracked[client_id].entity_id,
-                    tracked[client_id].client.mac,
-                )
-                tracked[client_id].async_schedule_update_ha_state()
                 continue
 
-            client = controller.api.clients[client_id]
+            tracked[item_id] = tracker_class(items[item_id], controller)
-
+            new_tracked.append(tracked[item_id])
-            if (
-                not client.is_wired
-                and controller.option_ssid_filter
-                and client.essid not in controller.option_ssid_filter
-            ):
-                continue
-
-            if not controller.option_track_wired_clients and client.is_wired:
-                continue
-
-            tracked[client_id] = UniFiClientTracker(client, controller)
-            new_tracked.append(tracked[client_id])
-            LOGGER.debug(
-                "New UniFi client tracker %s (%s)",
-                client.name or client.hostname,
-                client.mac,
-            )
-
-    if controller.option_track_devices:
-
-        for device_id in controller.api.devices:
-
-            if device_id in tracked:
-                if not tracked[device_id].enabled:
-                    continue
-                LOGGER.debug(
-                    "Updating UniFi tracked device %s (%s)",
-                    tracked[device_id].entity_id,
-                    tracked[device_id].device.mac,
-                )
-                tracked[device_id].async_schedule_update_ha_state()
-                continue
-
-            device = controller.api.devices[device_id]
-
-            tracked[device_id] = UniFiDeviceTracker(device, controller)
-            new_tracked.append(tracked[device_id])
-            LOGGER.debug(
-                "New UniFi device tracker %s (%s)",
-                device.name or device.model,
-                device.mac,
-            )
 
     if new_tracked:
         async_add_entities(new_tracked)
@@ -218,8 +191,33 @@ class UniFiClientTracker(ScannerEntity):
         self.client = client
         self.controller = controller
 
+    @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
+
+        if (
+            not self.client.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.client.is_wired:
+            return False
+
+        return True
+
+    async def async_added_to_hass(self):
+        """Client entity created."""
+        LOGGER.debug("New UniFi client tracker %s (%s)", self.name, self.client.mac)
+
     async def async_update(self):
         """Synchronize state with controller."""
+        LOGGER.debug(
+            "Updating UniFi tracked client %s (%s)", self.entity_id, self.client.mac
+        )
         await self.controller.request_update()
 
     @property
@@ -277,8 +275,23 @@ class UniFiDeviceTracker(ScannerEntity):
         self.device = device
         self.controller = controller
 
+    @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_devices:
+            return False
+
+        return True
+
+    async def async_added_to_hass(self):
+        """Subscribe to device events."""
+        LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac)
+
     async def async_update(self):
         """Synchronize state with controller."""
+        LOGGER.debug(
+            "Updating UniFi tracked device %s (%s)", self.entity_id, self.device.mac
+        )
         await self.controller.request_update()
 
     @property

homeassistant/components/unifi/switch.py

@@ -61,7 +61,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
         """Update the values of the controller."""
         update_items(controller, async_add_entities, switches, switches_off)
 
-    async_dispatcher_connect(hass, controller.event_update, update_controller)
+    async_dispatcher_connect(hass, controller.signal_update, update_controller)
 
     update_controller()
     switches_off.clear()
@@ -220,7 +220,7 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
             or self.client.sw_mac
             and (
                 self.controller.available
-                or self.client.sw_mac in self.controller.api.devices
+                and self.client.sw_mac in self.controller.api.devices
             )
         )
 

tests/components/unifi/test_controller.py

@@ -42,12 +42,12 @@ async def test_controller_setup():
             {
                 CONF_HOST: CONTROLLER_DATA[CONF_HOST],
                 CONF_SITE_ID: "nice name",
-                controller.CONF_BLOCK_CLIENT: [],
+                controller.CONF_BLOCK_CLIENT: ["mac"],
-                controller.CONF_TRACK_CLIENTS: True,
+                controller.CONF_DONT_TRACK_CLIENTS: True,
-                controller.CONF_TRACK_DEVICES: True,
+                controller.CONF_DONT_TRACK_DEVICES: True,
-                controller.CONF_TRACK_WIRED_CLIENTS: True,
+                controller.CONF_DONT_TRACK_WIRED_CLIENTS: True,
-                controller.CONF_DETECTION_TIME: 300,
+                controller.CONF_DETECTION_TIME: 30,
-                controller.CONF_SSID_FILTER: [],
+                controller.CONF_SSID_FILTER: ["ssid"],
             }
         ]
     }

tests/components/unifi/test_device_tracker.py

@@ -15,6 +15,8 @@ from homeassistant.components.unifi.const import (
     CONF_CONTROLLER,
     CONF_SITE_ID,
     CONF_SSID_FILTER,
+    CONF_TRACK_DEVICES,
+    CONF_TRACK_WIRED_CLIENTS,
     UNIFI_CONFIG,
 )
 from homeassistant.const import (
@@ -69,10 +71,10 @@ DEVICE_1 = {
     "mac": "00:00:00:00:01:01",
     "model": "US16P150",
     "name": "device_1",
-    "overheating": False,
+    "overheating": True,
     "state": 1,
     "type": "usw",
-    "upgradable": False,
+    "upgradable": True,
     "version": "4.0.42.10433",
 }
 DEVICE_2 = {
@@ -149,6 +151,7 @@ async def setup_controller(hass, mock_controller, options={}):
         system_options={},
         options=options,
     )
+    hass.config_entries._entries.append(config_entry)
     mock_controller.config_entry = config_entry
 
     await mock_controller.async_update()
@@ -230,6 +233,25 @@ async def test_tracked_devices(hass, mock_controller):
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1.state == STATE_UNAVAILABLE
 
+    mock_controller.config_entry.add_update_listener(
+        mock_controller.async_options_updated
+    )
+    hass.config_entries.async_update_entry(
+        mock_controller.config_entry,
+        options={
+            CONF_SSID_FILTER: [],
+            CONF_TRACK_WIRED_CLIENTS: False,
+            CONF_TRACK_DEVICES: False,
+        },
+    )
+    await hass.async_block_till_done()
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1
+    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 None
+
 
 async def test_restoring_client(hass, mock_controller):
     """Test the update_items function with some clients."""