From 39257164a92e13094ec37da0e1e1b6c50d1bd4c9 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 2 Aug 2019 10:13:00 +0200 Subject: [PATCH] UniFi - allow configuration to not track clients or devices (#25642) * Allow configuration to not track clients or devices --- homeassistant/components/unifi/__init__.py | 4 ++ homeassistant/components/unifi/const.py | 2 + .../components/unifi/device_tracker.py | 72 ++++++++++--------- tests/components/unifi/test_device_tracker.py | 39 +++++++++- 4 files changed, 84 insertions(+), 33 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 726d4793085..f4df139001d 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -11,6 +11,8 @@ from .const import ( CONF_BLOCK_CLIENT, CONF_CONTROLLER, CONF_DETECTION_TIME, + CONF_DONT_TRACK_CLIENTS, + CONF_DONT_TRACK_DEVICES, CONF_SITE_ID, CONF_SSID_FILTER, CONTROLLER_ID, @@ -28,6 +30,8 @@ CONTROLLER_SCHEMA = vol.Schema( vol.Optional(CONF_BLOCK_CLIENT, default=[]): vol.All( cv.ensure_list, [cv.string] ), + vol.Optional(CONF_DONT_TRACK_CLIENTS): cv.boolean, + vol.Optional(CONF_DONT_TRACK_DEVICES): cv.boolean, vol.Optional(CONF_DETECTION_TIME): vol.All( cv.time_period, cv.positive_timedelta ), diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index 383b018264a..1295849704c 100644 --- a/homeassistant/components/unifi/const.py +++ b/homeassistant/components/unifi/const.py @@ -13,6 +13,8 @@ UNIFI_CONFIG = "unifi_config" CONF_BLOCK_CLIENT = "block_client" CONF_DETECTION_TIME = "detection_time" +CONF_DONT_TRACK_CLIENTS = "dont_track_clients" +CONF_DONT_TRACK_DEVICES = "dont_track_devices" CONF_SSID_FILTER = "ssid_filter" ATTR_MANUFACTURER = "Ubiquiti Networks" diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 4046f5f63d2..ce5a1a7f608 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -28,6 +28,8 @@ from .const import ( ATTR_MANUFACTURER, CONF_CONTROLLER, CONF_DETECTION_TIME, + CONF_DONT_TRACK_CLIENTS, + CONF_DONT_TRACK_DEVICES, CONF_SITE_ID, CONF_SSID_FILTER, CONTROLLER_ID, @@ -154,46 +156,52 @@ def update_items(controller, async_add_entities, tracked): """Update tracked device state from the controller.""" new_tracked = [] - for client_id in controller.api.clients: + if not controller.unifi_config.get(CONF_DONT_TRACK_CLIENTS, False): - if client_id in tracked: + for client_id in controller.api.clients: + + if client_id in tracked: + 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] + + if ( + not client.is_wired + and CONF_SSID_FILTER in controller.unifi_config + and client.essid not in controller.unifi_config[CONF_SSID_FILTER] + ): + continue + + tracked[client_id] = UniFiClientTracker(client, controller) + new_tracked.append(tracked[client_id]) LOGGER.debug( - "Updating UniFi tracked client %s (%s)", - tracked[client_id].entity_id, - tracked[client_id].client.mac, + "New UniFi client tracker %s (%s)", client.hostname, client.mac ) - tracked[client_id].async_schedule_update_ha_state() - continue - client = controller.api.clients[client_id] + if not controller.unifi_config.get(CONF_DONT_TRACK_DEVICES, False): - if ( - not client.is_wired - and CONF_SSID_FILTER in controller.unifi_config - and client.essid not in controller.unifi_config[CONF_SSID_FILTER] - ): - continue + for device_id in controller.api.devices: - tracked[client_id] = UniFiClientTracker(client, controller) - new_tracked.append(tracked[client_id]) - LOGGER.debug("New UniFi client tracker %s (%s)", client.hostname, client.mac) + if device_id in tracked: + 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 - for device_id in controller.api.devices: + device = controller.api.devices[device_id] - if device_id in tracked: - 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, device.mac) + tracked[device_id] = UniFiDeviceTracker(device, controller) + new_tracked.append(tracked[device_id]) + LOGGER.debug("New UniFi device tracker %s (%s)", device.name, device.mac) if new_tracked: async_add_entities(new_tracked) diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 5accbb584b4..fb13bef42aa 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -22,6 +22,7 @@ from homeassistant.const import ( CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL, + STATE_UNAVAILABLE, ) from homeassistant.helpers import entity_registry from homeassistant.setup import async_setup_component @@ -211,7 +212,7 @@ async def test_tracked_devices(hass, mock_controller): await hass.async_block_till_done() device_1 = hass.states.get("device_tracker.device_1") - assert device_1.state == "unavailable" + assert device_1.state == STATE_UNAVAILABLE async def test_restoring_client(hass, mock_controller): @@ -243,3 +244,39 @@ async def test_restoring_client(hass, mock_controller): device_1 = hass.states.get("device_tracker.client_1") assert device_1 is not None + + +async def test_dont_track_clients(hass, mock_controller): + """Test dont track clients config works.""" + mock_controller.mock_client_responses.append([CLIENT_1]) + mock_controller.mock_device_responses.append([DEVICE_1]) + mock_controller.unifi_config = {unifi.CONF_DONT_TRACK_CLIENTS: True} + + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 2 + assert len(hass.states.async_all()) == 3 + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is None + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is not None + assert device_1.state == "not_home" + + +async def test_dont_track_devices(hass, mock_controller): + """Test dont track devices config works.""" + mock_controller.mock_client_responses.append([CLIENT_1]) + mock_controller.mock_device_responses.append([DEVICE_1]) + mock_controller.unifi_config = {unifi.CONF_DONT_TRACK_DEVICES: True} + + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 2 + assert len(hass.states.async_all()) == 3 + + client_1 = hass.states.get("device_tracker.client_1") + assert client_1 is not None + assert client_1.state == "not_home" + + device_1 = hass.states.get("device_tracker.device_1") + assert device_1 is None