UniFi - allow configuration to not track clients or devices (#25642)

* Allow configuration to not track clients or devices

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
         ),

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"

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)",
+                "New UniFi client tracker %s (%s)", client.hostname, client.mac
-                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 controller.unifi_config.get(CONF_DONT_TRACK_DEVICES, False):
 
-        if (
+        for device_id in controller.api.devices:
-            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)
+            if device_id in tracked:
-        new_tracked.append(tracked[client_id])
+                LOGGER.debug(
-        LOGGER.debug("New UniFi client tracker %s (%s)", client.hostname, client.mac)
+                    "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:
+            tracked[device_id] = UniFiDeviceTracker(device, controller)
-            LOGGER.debug(
+            new_tracked.append(tracked[device_id])
-                "Updating UniFi tracked device %s (%s)",
+            LOGGER.debug("New UniFi device tracker %s (%s)", device.name, device.mac)
-                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)
 
     if new_tracked:
         async_add_entities(new_tracked)

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