From a6e7027ef8fcd4569cd6f25fefd56b3d58200c24 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Thu, 1 Aug 2019 23:04:17 +0200 Subject: [PATCH 1/3] Allow configuration to not track clients or devices --- homeassistant/components/unifi/__init__.py | 4 ++ homeassistant/components/unifi/const.py | 2 + homeassistant/components/unifi/controller.py | 12 ++++ .../components/unifi/device_tracker.py | 68 ++++++++++--------- tests/components/unifi/test_device_tracker.py | 39 ++++++++++- 5 files changed, 93 insertions(+), 32 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 726d479308516..f4df139001d0f 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -11,6 +11,8 @@ 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 @@ 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 383b018264a80..1295849704c3f 100644 --- a/homeassistant/components/unifi/const.py +++ b/homeassistant/components/unifi/const.py @@ -13,6 +13,8 @@ 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/controller.py b/homeassistant/components/unifi/controller.py index cb82e6cf1c1d0..e6e242ccdd4f0 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -15,6 +15,8 @@ from .const import ( CONF_BLOCK_CLIENT, CONF_CONTROLLER, + CONF_DONT_TRACK_CLIENTS, + CONF_DONT_TRACK_DEVICES, CONF_SITE_ID, CONTROLLER_ID, LOGGER, @@ -63,6 +65,16 @@ def block_clients(self): """Return list of clients to block.""" return self.unifi_config.get(CONF_BLOCK_CLIENT, []) + @property + def dont_track_clients(self): + """Return if clients should be tracked.""" + return self.unifi_config.get(CONF_DONT_TRACK_CLIENTS, False) + + @property + def dont_track_devices(self): + """Return if devices should be tracked.""" + return self.unifi_config.get(CONF_DONT_TRACK_DEVICES, False) + @property def mac(self): """Return the mac address of this controller.""" diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 4046f5f63d223..607bfb04d1c82 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -154,46 +154,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.dont_track_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] + for client_id in controller.api.clients: - 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 + 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 - tracked[client_id] = UniFiClientTracker(client, controller) - new_tracked.append(tracked[client_id]) - LOGGER.debug("New UniFi client tracker %s (%s)", client.hostname, client.mac) + client = controller.api.clients[client_id] - for device_id in controller.api.devices: + 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 - if device_id in tracked: + tracked[client_id] = UniFiClientTracker(client, controller) + new_tracked.append(tracked[client_id]) LOGGER.debug( - "Updating UniFi tracked device %s (%s)", - tracked[device_id].entity_id, - tracked[device_id].device.mac, + "New UniFi client tracker %s (%s)", client.hostname, client.mac ) - tracked[device_id].async_schedule_update_ha_state() - continue - device = controller.api.devices[device_id] + if not controller.dont_track_devices: + + for device_id in controller.api.devices: + + 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 5accbb584b48a..114c25ad2c8f5 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -22,6 +22,7 @@ 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 the update_items function with some clients.""" + 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 the update_items function with some clients.""" + 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 From 9e21a88143df2726a3400e8a3ab16ae2094fba26 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 2 Aug 2019 08:37:01 +0200 Subject: [PATCH 2/3] Fix Andrews comments --- homeassistant/components/unifi/controller.py | 10 ---------- homeassistant/components/unifi/device_tracker.py | 6 ++++-- tests/components/unifi/test_device_tracker.py | 4 ++-- 3 files changed, 6 insertions(+), 14 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index e6e242ccdd4f0..01b073da7257e 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -65,16 +65,6 @@ def block_clients(self): """Return list of clients to block.""" return self.unifi_config.get(CONF_BLOCK_CLIENT, []) - @property - def dont_track_clients(self): - """Return if clients should be tracked.""" - return self.unifi_config.get(CONF_DONT_TRACK_CLIENTS, False) - - @property - def dont_track_devices(self): - """Return if devices should be tracked.""" - return self.unifi_config.get(CONF_DONT_TRACK_DEVICES, False) - @property def mac(self): """Return the mac address of this controller.""" diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 607bfb04d1c82..ce5a1a7f60867 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -28,6 +28,8 @@ ATTR_MANUFACTURER, CONF_CONTROLLER, CONF_DETECTION_TIME, + CONF_DONT_TRACK_CLIENTS, + CONF_DONT_TRACK_DEVICES, CONF_SITE_ID, CONF_SSID_FILTER, CONTROLLER_ID, @@ -154,7 +156,7 @@ def update_items(controller, async_add_entities, tracked): """Update tracked device state from the controller.""" new_tracked = [] - if not controller.dont_track_clients: + if not controller.unifi_config.get(CONF_DONT_TRACK_CLIENTS, False): for client_id in controller.api.clients: @@ -182,7 +184,7 @@ def update_items(controller, async_add_entities, tracked): "New UniFi client tracker %s (%s)", client.hostname, client.mac ) - if not controller.dont_track_devices: + if not controller.unifi_config.get(CONF_DONT_TRACK_DEVICES, False): for device_id in controller.api.devices: diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 114c25ad2c8f5..fb13bef42aa53 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -247,7 +247,7 @@ async def test_restoring_client(hass, mock_controller): async def test_dont_track_clients(hass, mock_controller): - """Test the update_items function with some clients.""" + """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} @@ -265,7 +265,7 @@ async def test_dont_track_clients(hass, mock_controller): async def test_dont_track_devices(hass, mock_controller): - """Test the update_items function with some clients.""" + """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} From 21d51e44458d73201792ed1aabf5abc54dce836e Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 2 Aug 2019 08:37:25 +0200 Subject: [PATCH 3/3] Remove unnecessary import --- homeassistant/components/unifi/controller.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 01b073da7257e..cb82e6cf1c1d0 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -15,8 +15,6 @@ from .const import ( CONF_BLOCK_CLIENT, CONF_CONTROLLER, - CONF_DONT_TRACK_CLIENTS, - CONF_DONT_TRACK_DEVICES, CONF_SITE_ID, CONTROLLER_ID, LOGGER,