diff --git a/homeassistant/components/unifi/.translations/en.json b/homeassistant/components/unifi/.translations/en.json index d42a647c82fe2c..881a80ff65e128 100644 --- a/homeassistant/components/unifi/.translations/en.json +++ b/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", diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 781fcdeae828e8..26c46950e8e38f 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/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), diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index 4acf75fbdd9f56..770a2d9a5bde8a 100644 --- a/homeassistant/components/unifi/const.py +++ b/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" diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 864d131d287f31..32002d84c55d46 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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,32 +90,25 @@ 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.""" @@ -122,6 +116,16 @@ def option_track_wired_clients(self): 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): diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index c5aa74706a1cbd..a76da481356e76 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/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): diff --git a/homeassistant/components/unifi/strings.json b/homeassistant/components/unifi/strings.json index 881e97bc9cad72..65f22a8af436cc 100644 --- a/homeassistant/components/unifi/strings.json +++ b/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" } -} +} \ No newline at end of file diff --git a/homeassistant/components/unifi/unifi_client.py b/homeassistant/components/unifi/unifi_client.py index 644b0856bb49dd..11b23af720d0b7 100644 --- a/homeassistant/components/unifi/unifi_client.py +++ b/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: diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index 3366ec1641d101..d7db618f2c2916 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/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, } diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index cfb4637a6c4fda..238a68dd4b4e6f 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/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"