Rework UniFi client configuration (#99483)

homeassistant/components/unifi/config_flow.py

@@ -34,6 +34,7 @@ from .const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_ALLOW_UPTIME_SENSORS,
     CONF_BLOCK_CLIENT,
+    CONF_CLIENT_SOURCE,
     CONF_DETECTION_TIME,
     CONF_DPI_RESTRICTIONS,
     CONF_IGNORE_WIRED_BUG,
@@ -257,7 +258,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
         self.options[CONF_BLOCK_CLIENT] = self.controller.option_block_clients
 
         if self.show_advanced_options:
-            return await self.async_step_device_tracker()
+            return await self.async_step_configure_entity_sources()
 
         return await self.async_step_simple_options()
 
@@ -296,6 +297,32 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
             last_step=True,
         )
 
+    async def async_step_configure_entity_sources(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Select sources for entities."""
+        if user_input is not None:
+            self.options.update(user_input)
+            return await self.async_step_device_tracker()
+
+        clients = {
+            client.mac: f"{client.name or client.hostname} ({client.mac})"
+            for client in self.controller.api.clients.values()
+        }
+
+        return self.async_show_form(
+            step_id="configure_entity_sources",
+            data_schema=vol.Schema(
+                {
+                    vol.Optional(
+                        CONF_CLIENT_SOURCE,
+                        default=self.options.get(CONF_CLIENT_SOURCE, []),
+                    ): cv.multi_select(clients),
+                }
+            ),
+            last_step=False,
+        )
+
     async def async_step_device_tracker(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:

homeassistant/components/unifi/const.py

@@ -23,6 +23,7 @@ UNIFI_WIRELESS_CLIENTS = "unifi_wireless_clients"
 CONF_ALLOW_BANDWIDTH_SENSORS = "allow_bandwidth_sensors"
 CONF_ALLOW_UPTIME_SENSORS = "allow_uptime_sensors"
 CONF_BLOCK_CLIENT = "block_client"
+CONF_CLIENT_SOURCE = "client_source"
 CONF_DETECTION_TIME = "detection_time"
 CONF_DPI_RESTRICTIONS = "dpi_restrictions"
 CONF_IGNORE_WIRED_BUG = "ignore_wired_bug"

homeassistant/components/unifi/controller.py

@@ -47,6 +47,7 @@ from .const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_ALLOW_UPTIME_SENSORS,
     CONF_BLOCK_CLIENT,
+    CONF_CLIENT_SOURCE,
     CONF_DETECTION_TIME,
     CONF_DPI_RESTRICTIONS,
     CONF_IGNORE_WIRED_BUG,
@@ -109,6 +110,9 @@ class UniFiController:
         """Store attributes to avoid property call overhead since they are called frequently."""
         options = self.config_entry.options
 
+        # Allow creating entities from clients.
+        self.option_supported_clients: list[str] = options.get(CONF_CLIENT_SOURCE, [])
+
         # Device tracker options
 
         # Config entry option to not track clients.

homeassistant/components/unifi/device_tracker.py

@@ -80,6 +80,9 @@ WIRELESS_DISCONNECTION = (
 @callback
 def async_client_allowed_fn(controller: UniFiController, obj_id: str) -> bool:
     """Check if client is allowed."""
+    if obj_id in controller.option_supported_clients:
+        return True
+
     if not controller.option_track_clients:
         return False
 

homeassistant/components/unifi/sensor.py

@@ -49,6 +49,22 @@ from .entity import (
 )
 
 
+@callback
+def async_bandwidth_sensor_allowed_fn(controller: UniFiController, obj_id: str) -> bool:
+    """Check if client is allowed."""
+    if obj_id in controller.option_supported_clients:
+        return True
+    return controller.option_allow_bandwidth_sensors
+
+
+@callback
+def async_uptime_sensor_allowed_fn(controller: UniFiController, obj_id: str) -> bool:
+    """Check if client is allowed."""
+    if obj_id in controller.option_supported_clients:
+        return True
+    return controller.option_allow_uptime_sensors
+
+
 @callback
 def async_client_rx_value_fn(controller: UniFiController, client: Client) -> float:
     """Calculate receiving data transfer value."""
@@ -139,7 +155,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         native_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND,
         icon="mdi:upload",
         has_entity_name=True,
-        allowed_fn=lambda controller, _: controller.option_allow_bandwidth_sensors,
+        allowed_fn=async_bandwidth_sensor_allowed_fn,
         api_handler_fn=lambda api: api.clients,
         available_fn=lambda controller, _: controller.available,
         device_info_fn=async_client_device_info_fn,
@@ -159,7 +175,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         native_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND,
         icon="mdi:download",
         has_entity_name=True,
-        allowed_fn=lambda controller, _: controller.option_allow_bandwidth_sensors,
+        allowed_fn=async_bandwidth_sensor_allowed_fn,
         api_handler_fn=lambda api: api.clients,
         available_fn=lambda controller, _: controller.available,
         device_info_fn=async_client_device_info_fn,
@@ -198,7 +214,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
         entity_category=EntityCategory.DIAGNOSTIC,
         has_entity_name=True,
         entity_registry_enabled_default=False,
-        allowed_fn=lambda controller, _: controller.option_allow_uptime_sensors,
+        allowed_fn=async_uptime_sensor_allowed_fn,
         api_handler_fn=lambda api: api.clients,
         available_fn=lambda controller, obj_id: controller.available,
         device_info_fn=async_client_device_info_fn,

homeassistant/components/unifi/strings.json

@@ -30,6 +30,13 @@
       "integration_not_setup": "UniFi integration is not set up"
     },
     "step": {
+      "configure_entity_sources": {
+        "data": {
+          "client_source": "Create entities from network clients"
+        },
+        "description": "Select sources to create entities from",
+        "title": "UniFi Network Entity Sources"
+      },
       "device_tracker": {
         "data": {
           "detection_time": "Time in seconds from last seen until considered away",

homeassistant/components/unifi/switch.py

@@ -60,6 +60,14 @@ CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKE
 CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED)
 
 
+@callback
+def async_block_client_allowed_fn(controller: UniFiController, obj_id: str) -> bool:
+    """Check if client is allowed."""
+    if obj_id in controller.option_supported_clients:
+        return True
+    return obj_id in controller.option_block_clients
+
+
 @callback
 def async_dpi_group_is_on_fn(
     controller: UniFiController, dpi_group: DPIRestrictionGroup
@@ -198,7 +206,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
         entity_category=EntityCategory.CONFIG,
         has_entity_name=True,
         icon="mdi:ethernet",
-        allowed_fn=lambda controller, obj_id: obj_id in controller.option_block_clients,
+        allowed_fn=async_block_client_allowed_fn,
         api_handler_fn=lambda api: api.clients,
         available_fn=lambda controller, obj_id: controller.available,
         control_fn=async_block_client_control_fn,

tests/components/unifi/test_config_flow.py

@@ -11,6 +11,7 @@ from homeassistant.components.unifi.const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
     CONF_ALLOW_UPTIME_SENSORS,
     CONF_BLOCK_CLIENT,
+    CONF_CLIENT_SOURCE,
     CONF_DETECTION_TIME,
     CONF_DPI_RESTRICTIONS,
     CONF_IGNORE_WIRED_BUG,
@@ -462,6 +463,17 @@ async def test_advanced_option_flow(
         config_entry.entry_id, context={"show_advanced_options": True}
     )
 
+    assert result["type"] == data_entry_flow.FlowResultType.FORM
+    assert result["step_id"] == "configure_entity_sources"
+    assert not result["last_step"]
+    assert list(result["data_schema"].schema[CONF_CLIENT_SOURCE].options.keys()) == [
+        "00:00:00:00:00:01"
+    ]
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"],
+        user_input={CONF_CLIENT_SOURCE: ["00:00:00:00:00:01"]},
+    )
+
     assert result["type"] == data_entry_flow.FlowResultType.FORM
     assert result["step_id"] == "device_tracker"
     assert not result["last_step"]
@@ -510,6 +522,7 @@ async def test_advanced_option_flow(
 
     assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
     assert result["data"] == {
+        CONF_CLIENT_SOURCE: ["00:00:00:00:00:01"],
         CONF_TRACK_CLIENTS: False,
         CONF_TRACK_WIRED_CLIENTS: False,
         CONF_TRACK_DEVICES: False,

tests/components/unifi/test_device_tracker.py

@@ -9,6 +9,7 @@ from homeassistant import config_entries
 from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
 from homeassistant.components.unifi.const import (
     CONF_BLOCK_CLIENT,
+    CONF_CLIENT_SOURCE,
     CONF_IGNORE_WIRED_BUG,
     CONF_SSID_FILTER,
     CONF_TRACK_CLIENTS,
@@ -132,21 +133,29 @@ async def test_tracked_clients(
         "last_seen": None,
         "mac": "00:00:00:00:00:05",
     }
+    client_6 = {
+        "hostname": "client_6",
+        "ip": "10.0.0.6",
+        "is_wired": True,
+        "last_seen": 1562600145,
+        "mac": "00:00:00:00:00:06",
+    }
 
     await setup_unifi_integration(
         hass,
         aioclient_mock,
-        options={CONF_SSID_FILTER: ["ssid"]},
+        options={CONF_SSID_FILTER: ["ssid"], CONF_CLIENT_SOURCE: [client_6["mac"]]},
-        clients_response=[client_1, client_2, client_3, client_4, client_5],
+        clients_response=[client_1, client_2, client_3, client_4, client_5, client_6],
         known_wireless_clients=(client_4["mac"],),
     )
 
-    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 4
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 5
     assert hass.states.get("device_tracker.client_1").state == STATE_NOT_HOME
     assert hass.states.get("device_tracker.client_2").state == STATE_NOT_HOME
     assert (
         hass.states.get("device_tracker.client_5").attributes["host_name"] == "client_5"
     )
+    assert hass.states.get("device_tracker.client_6").state == STATE_NOT_HOME
 
     # Client on SSID not in SSID filter
     assert not hass.states.get("device_tracker.client_3")