From 721c45b7a318de1a05c79c5fb0940b3885ac59fa Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Sun, 22 Oct 2023 23:39:54 +0200 Subject: [PATCH] Rework UniFi client configuration (#99483) --- homeassistant/components/unifi/config_flow.py | 29 ++++++++++++++++++- homeassistant/components/unifi/const.py | 1 + homeassistant/components/unifi/controller.py | 4 +++ .../components/unifi/device_tracker.py | 3 ++ homeassistant/components/unifi/sensor.py | 22 ++++++++++++-- homeassistant/components/unifi/strings.json | 7 +++++ homeassistant/components/unifi/switch.py | 10 ++++++- tests/components/unifi/test_config_flow.py | 13 +++++++++ tests/components/unifi/test_device_tracker.py | 15 ++++++++-- 9 files changed, 96 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 8c0696463c5..a678517eca9 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/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: diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index 176511645aa..c78313f66e2 100644 --- a/homeassistant/components/unifi/const.py +++ b/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" diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index b0ce43fe959..108ff87d026 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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. diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 22a530e0369..5c9694c669c 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/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 diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index f07269dfdd2..3d0ffa1896e 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/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, diff --git a/homeassistant/components/unifi/strings.json b/homeassistant/components/unifi/strings.json index e441d4695ed..9c609ca8c07 100644 --- a/homeassistant/components/unifi/strings.json +++ b/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", diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 0aa39914686..41c1f55a22a 100644 --- a/homeassistant/components/unifi/switch.py +++ b/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, diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index ba906f86eef..71d2dc038c1 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/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, diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 2680a357d77..cbff868d9a6 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/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"]}, - clients_response=[client_1, client_2, client_3, client_4, client_5], + options={CONF_SSID_FILTER: ["ssid"], CONF_CLIENT_SOURCE: [client_6["mac"]]}, + 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")