UniFi - Allow user to configure tracking of specific clients

homeassistant/components/unifi/.translations/en.json

@@ -40,6 +40,7 @@
             },
             "device_tracker": {
                 "data": {
+                    "clients_to_track": "Select specific clients to track",
                     "detection_time": "Time in seconds from last seen until considered away",
                     "ssid_filter": "Select SSIDs to track wireless clients on",
                     "track_clients": "Track network clients",

homeassistant/components/unifi/config_flow.py

@@ -17,6 +17,7 @@
 from .const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_BLOCK_CLIENT,
+    CONF_CLIENTS_TO_TRACK,
     CONF_CONTROLLER,
     CONF_DETECTION_TIME,
     CONF_POE_CLIENTS,
@@ -189,6 +190,11 @@ async def async_step_device_tracker(self, user_input=None):
             self.options.update(user_input)
             return await self.async_step_client_control()
 
+        clients_to_track = {
+            client.mac: client.name or client.hostname
+            for client in self.controller.api.clients.values()
+        }
+
         ssid_filter = {wlan: wlan for wlan in self.controller.api.wlans}
 
         return self.async_show_form(
@@ -207,6 +213,10 @@ async def async_step_device_tracker(self, user_input=None):
                         CONF_TRACK_DEVICES,
                         default=self.controller.option_track_devices,
                     ): bool,
+                    vol.Optional(
+                        CONF_CLIENTS_TO_TRACK,
+                        default=self.controller.option_clients_to_track,
+                    ): cv.multi_select(clients_to_track),
                     vol.Optional(
                         CONF_SSID_FILTER, default=self.controller.option_ssid_filter
                     ): cv.multi_select(ssid_filter),

homeassistant/components/unifi/const.py

@@ -13,6 +13,7 @@
 
 CONF_ALLOW_BANDWIDTH_SENSORS = "allow_bandwidth_sensors"
 CONF_BLOCK_CLIENT = "block_client"
+CONF_CLIENTS_TO_TRACK = "clients_to_track"
 CONF_DETECTION_TIME = "detection_time"
 CONF_POE_CLIENTS = "poe_clients"
 CONF_TRACK_CLIENTS = "track_clients"

homeassistant/components/unifi/controller.py

@@ -22,6 +22,7 @@
 from .const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_BLOCK_CLIENT,
+    CONF_CLIENTS_TO_TRACK,
     CONF_CONTROLLER,
     CONF_DETECTION_TIME,
     CONF_POE_CLIENTS,
@@ -89,39 +90,42 @@ def site_role(self):
         return self._site_role
 
     @property
-    def option_allow_bandwidth_sensors(self):
-        """Config entry option to allow bandwidth sensors."""
-        return self.config_entry.options.get(
-            CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS
-        )
+    def mac(self):
+        """Return the mac address of this controller."""
+        for client in self.api.clients.values():
+            if self.host == client.ip:
+                return client.mac
+        return None
 
-    @property
-    def option_block_clients(self):
-        """Config entry option with list of clients to control network access."""
-        return self.config_entry.options.get(CONF_BLOCK_CLIENT, [])
+    # Device tracker options
 
     @property
-    def option_poe_clients(self):
-        """Config entry option to control poe clients."""
-        return self.config_entry.options.get(CONF_POE_CLIENTS, DEFAULT_POE_CLIENTS)
+    def option_clients_to_track(self):
+        """Config entry option with list of clients to control network access."""
+        return self.config_entry.options.get(CONF_CLIENTS_TO_TRACK, [])
 
     @property
     def option_track_clients(self):
         """Config entry option to not track clients."""
         return self.config_entry.options.get(CONF_TRACK_CLIENTS, DEFAULT_TRACK_CLIENTS)
 
-    @property
-    def option_track_devices(self):
-        """Config entry option to not track devices."""
-        return self.config_entry.options.get(CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES)
-
     @property
     def option_track_wired_clients(self):
         """Config entry option to not track wired clients."""
         return self.config_entry.options.get(
             CONF_TRACK_WIRED_CLIENTS, DEFAULT_TRACK_WIRED_CLIENTS
         )
 
+    @property
+    def option_track_devices(self):
+        """Config entry option to not track devices."""
+        return self.config_entry.options.get(CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES)
+
+    @property
+    def option_ssid_filter(self):
+        """Config entry option listing what SSIDs are being used to track clients."""
+        return self.config_entry.options.get(CONF_SSID_FILTER, [])
+
     @property
     def option_detection_time(self):
         """Config entry option defining number of seconds from last seen to away."""
@@ -131,18 +135,26 @@ def option_detection_time(self):
             )
         )
 
+    # Client control options
+
     @property
-    def option_ssid_filter(self):
-        """Config entry option listing what SSIDs are being used to track clients."""
-        return self.config_entry.options.get(CONF_SSID_FILTER, [])
+    def option_poe_clients(self):
+        """Config entry option to control poe clients."""
+        return self.config_entry.options.get(CONF_POE_CLIENTS, DEFAULT_POE_CLIENTS)
 
     @property
-    def mac(self):
-        """Return the mac address of this controller."""
-        for client in self.api.clients.values():
-            if self.host == client.ip:
-                return client.mac
-        return None
+    def option_block_clients(self):
+        """Config entry option with list of clients to control network access."""
+        return self.config_entry.options.get(CONF_BLOCK_CLIENT, [])
+
+    # Statistics sensors options
+
+    @property
+    def option_allow_bandwidth_sensors(self):
+        """Config entry option to allow bandwidth sensors."""
+        return self.config_entry.options.get(
+            CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS
+        )
 
     @callback
     def async_unifi_signalling_callback(self, signal, data):

homeassistant/components/unifi/device_tracker.py

@@ -47,6 +47,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     option_track_clients = controller.option_track_clients
     option_track_devices = controller.option_track_devices
     option_track_wired_clients = controller.option_track_wired_clients
+    option_clients_to_track = controller.option_clients_to_track
     option_ssid_filter = controller.option_ssid_filter
 
     registry = await hass.helpers.entity_registry.async_get_registry()
@@ -89,6 +90,7 @@ def options_updated():
         nonlocal option_track_clients
         nonlocal option_track_devices
         nonlocal option_track_wired_clients
+        nonlocal option_clients_to_track
         nonlocal option_ssid_filter
 
         update = False
@@ -120,21 +122,34 @@ def options_updated():
                     if isinstance(entity, UniFiClientTracker) and entity.is_wired:
                         remove.add(mac)
 
+        if option_clients_to_track != controller.option_clients_to_track:
+            update = True
+
+            if controller.option_clients_to_track:
+                for mac, entity in tracked.items():
+                    if (
+                        isinstance(entity, UniFiClientTracker)
+                        and mac not in controller.option_clients_to_track
+                    ):
+                        remove.add(mac)
+
         if option_ssid_filter != controller.option_ssid_filter:
-            option_ssid_filter = controller.option_ssid_filter
             update = True
 
-            for mac, entity in tracked.items():
-                if (
-                    isinstance(entity, UniFiClientTracker)
-                    and not entity.is_wired
-                    and entity.client.essid not in option_ssid_filter
-                ):
-                    remove.add(mac)
+            if controller.option_ssid_filter:
+                for mac, entity in tracked.items():
+                    if (
+                        isinstance(entity, UniFiClientTracker)
+                        and not entity.is_wired
+                        and entity.client.essid not in controller.option_ssid_filter
+                    ):
+                        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
+        option_clients_to_track = controller.option_clients_to_track
+        option_ssid_filter = controller.option_ssid_filter
 
         for mac in remove:
             entity = tracked.pop(mac)
@@ -176,6 +191,12 @@ def add_entities(controller, async_add_entities, tracked):
             if tracker_class is UniFiClientTracker:
                 client = items[item_id]
 
+                if (
+                    controller.option_clients_to_track
+                    and client.mac not in controller.option_clients_to_track
+                ):
+                    continue
+
                 if not controller.option_track_wired_clients and client.is_wired:
                     continue
 
@@ -304,13 +325,12 @@ def __init__(self, device, controller):
         """Set up tracked device."""
         self.device = device
         self.controller = controller
-        self.listeners = []
 
     async def async_added_to_hass(self):
         """Subscribe to device events."""
         LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac)
         self.device.register_callback(self.async_update_callback)
-        self.listeners.append(
+        self.async_on_remove(
             async_dispatcher_connect(
                 self.hass, self.controller.signal_reachable, self.async_update_callback
             )
@@ -319,8 +339,6 @@ async def async_added_to_hass(self):
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect device object when removed."""
         self.device.remove_callback(self.async_update_callback)
-        for unsub_dispatcher in self.listeners:
-            unsub_dispatcher()
 
     @callback
     def async_update_callback(self):

homeassistant/components/unifi/strings.json

@@ -31,6 +31,7 @@
       },
       "device_tracker": {
         "data": {
+          "clients_to_track": "Select specific clients to track",
           "detection_time": "Time in seconds from last seen until considered away",
           "ssid_filter": "Select SSIDs to track wireless clients on",
           "track_clients": "Track network clients",
@@ -61,4 +62,4 @@
   "error": {
     "unknown_client_mac": "No client available in UniFi on that MAC address"
   }
-}
+}

homeassistant/components/unifi/unifi_client.py

@@ -39,7 +39,6 @@ def __init__(self, client, controller) -> None:
         """Set up client."""
         self.client = client
         self.controller = controller
-        self.listeners = []
 
         self.is_wired = self.client.mac not in controller.wireless_clients
         self.is_blocked = self.client.blocked
@@ -50,7 +49,7 @@ async def async_added_to_hass(self) -> None:
         """Client entity created."""
         LOGGER.debug("New UniFi client %s (%s)", self.name, self.client.mac)
         self.client.register_callback(self.async_update_callback)
-        self.listeners.append(
+        self.async_on_remove(
             async_dispatcher_connect(
                 self.hass, self.controller.signal_reachable, self.async_update_callback
             )
@@ -59,8 +58,6 @@ async def async_added_to_hass(self) -> None:
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect client object when removed."""
         self.client.remove_callback(self.async_update_callback)
-        for unsub_dispatcher in self.listeners:
-            unsub_dispatcher()
 
     @callback
     def async_update_callback(self) -> None:

tests/components/unifi/test_config_flow.py

@@ -9,6 +9,7 @@
 from homeassistant.components.unifi.const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_BLOCK_CLIENT,
+    CONF_CLIENTS_TO_TRACK,
     CONF_CONTROLLER,
     CONF_DETECTION_TIME,
     CONF_POE_CLIENTS,
@@ -284,6 +285,7 @@ async def test_option_flow(hass):
             CONF_TRACK_CLIENTS: False,
             CONF_TRACK_WIRED_CLIENTS: False,
             CONF_TRACK_DEVICES: False,
+            CONF_CLIENTS_TO_TRACK: ["00:00:00:00:00:01"],
             CONF_SSID_FILTER: ["SSID 1"],
             CONF_DETECTION_TIME: 100,
         },
@@ -338,9 +340,10 @@ async def test_option_flow(hass):
         CONF_TRACK_CLIENTS: False,
         CONF_TRACK_WIRED_CLIENTS: False,
         CONF_TRACK_DEVICES: False,
-        CONF_DETECTION_TIME: 100,
+        CONF_CLIENTS_TO_TRACK: ["00:00:00:00:00:01"],
         CONF_SSID_FILTER: ["SSID 1"],
-        CONF_BLOCK_CLIENT: ["00:00:00:00:00:01"],
+        CONF_DETECTION_TIME: 100,
         CONF_POE_CLIENTS: False,
+        CONF_BLOCK_CLIENT: ["00:00:00:00:00:01"],
         CONF_ALLOW_BANDWIDTH_SENSORS: True,
     }

tests/components/unifi/test_device_tracker.py

@@ -11,6 +11,7 @@
 import homeassistant.components.device_tracker as device_tracker
 from homeassistant.components.unifi.const import (
     CONF_BLOCK_CLIENT,
+    CONF_CLIENTS_TO_TRACK,
     CONF_SSID_FILTER,
     CONF_TRACK_CLIENTS,
     CONF_TRACK_DEVICES,
@@ -345,14 +346,59 @@ async def test_option_track_devices(hass):
     assert device_1 is not None
 
 
+async def test_option_clients_to_track(hass):
+    """Test the clients to track filter works."""
+    controller = await setup_unifi_integration(
+        hass, clients_response=[CLIENT_1, CLIENT_2],
+    )
+    assert len(hass.states.async_entity_ids("device_tracker")) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1
+    wired_client = hass.states.get("device_tracker.wired_client")
+    assert wired_client
+
+    # Set client filter to track only client_1
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_CLIENTS_TO_TRACK: [CLIENT_1["mac"]]},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids("device_tracker")) == 1
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1
+    wired_client = hass.states.get("device_tracker.wired_client")
+    assert not wired_client
+
+    # Clear client filter
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_CLIENTS_TO_TRACK: []},
+    )
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids("device_tracker")) == 2
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1
+    wired_client = hass.states.get("device_tracker.wired_client")
+    assert wired_client
+
+
 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],
+    controller = await setup_unifi_integration(hass, clients_response=[CLIENT_3],)
+    assert len(hass.states.async_entity_ids("device_tracker")) == 1
+
+    client_3 = hass.states.get("device_tracker.client_3")
+    assert client_3
+
+    # Set SSID filter
+    hass.config_entries.async_update_entry(
+        controller.config_entry, options={CONF_SSID_FILTER: ["ssid"]},
     )
-    assert len(hass.states.async_entity_ids("device_tracker")) == 0
+    await hass.async_block_till_done()
 
-    # SSID filter active
     client_3 = hass.states.get("device_tracker.client_3")
     assert not client_3
 
@@ -374,7 +420,6 @@ async def test_option_ssid_filter(hass):
     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"