From 9ff45ca01340804da331bb5629a5e17d4e2cce26 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Mon, 6 Mar 2023 16:08:14 +0100 Subject: [PATCH] Allow loading UniFi entities on config options change (#88762) Co-authored-by: Franck Nijhof --- homeassistant/components/unifi/controller.py | 18 +++++- homeassistant/components/unifi/entity.py | 9 +++ tests/components/unifi/test_device_tracker.py | 21 ++++++- tests/components/unifi/test_sensor.py | 62 +++++++++++++++++++ 4 files changed, 108 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 31d07278920..69c3cc78059 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -31,7 +31,10 @@ from homeassistant.helpers import ( entity_registry as er, ) from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC -from homeassistant.helpers.dispatcher import async_dispatcher_send +from homeassistant.helpers.dispatcher import ( + async_dispatcher_connect, + async_dispatcher_send, +) from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_registry import async_entries_for_config_entry from homeassistant.helpers.event import async_track_time_interval @@ -108,6 +111,7 @@ class UniFiController: self.load_config_entry_options() self.entities = {} + self.known_objects: set[tuple[str, str]] = set() def load_config_entry_options(self): """Store attributes to avoid property call overhead since they are called frequently.""" @@ -207,6 +211,7 @@ class UniFiController: [ unifi_platform_entity(obj_id, self, description) for obj_id in obj_ids + if (description.key, obj_id) not in self.known_objects if description.allowed_fn(self, obj_id) if description.supported_fn(self, obj_id) ] @@ -221,6 +226,17 @@ class UniFiController: api_handler.subscribe(async_create_entity, ItemEvent.ADDED) + @callback + def async_options_updated() -> None: + """Load new entities based on changed options.""" + async_add_unifi_entity(list(api_handler)) + + self.config_entry.async_on_unload( + async_dispatcher_connect( + self.hass, self.signal_options_update, async_options_updated + ) + ) + for description in descriptions: async_load_entities(description) diff --git a/homeassistant/components/unifi/entity.py b/homeassistant/components/unifi/entity.py index 783950310e4..5d763ecfe8a 100644 --- a/homeassistant/components/unifi/entity.py +++ b/homeassistant/components/unifi/entity.py @@ -103,6 +103,8 @@ class UnifiEntity(Entity, Generic[HandlerT, DataT]): self.controller = controller self.entity_description = description + controller.known_objects.add((description.key, obj_id)) + self._removed = False self._attr_available = description.available_fn(controller, obj_id) @@ -118,6 +120,13 @@ class UnifiEntity(Entity, Generic[HandlerT, DataT]): description = self.entity_description handler = description.api_handler_fn(self.controller.api) + @callback + def unregister_object() -> None: + """Remove object ID from known_objects when unloaded.""" + self.controller.known_objects.discard((description.key, self._obj_id)) + + self.async_on_remove(unregister_object) + # New data from handler self.async_on_remove( handler.subscribe( diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index f271394df1c..5dcf1fc6932 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -635,6 +635,15 @@ async def test_option_track_devices( assert hass.states.get("device_tracker.client") assert not hass.states.get("device_tracker.device") + hass.config_entries.async_update_entry( + config_entry, + options={CONF_TRACK_DEVICES: True}, + ) + await hass.async_block_till_done() + + assert hass.states.get("device_tracker.client") + assert hass.states.get("device_tracker.device") + async def test_option_ssid_filter( hass: HomeAssistant, @@ -1041,7 +1050,7 @@ async def test_dont_track_devices( "version": "4.0.42.10433", } - await setup_unifi_integration( + config_entry = await setup_unifi_integration( hass, aioclient_mock, options={CONF_TRACK_DEVICES: False}, @@ -1053,6 +1062,16 @@ async def test_dont_track_devices( assert hass.states.get("device_tracker.client") assert not hass.states.get("device_tracker.device") + hass.config_entries.async_update_entry( + config_entry, + options={CONF_TRACK_DEVICES: True}, + ) + await hass.async_block_till_done() + + assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2 + assert hass.states.get("device_tracker.client") + assert hass.states.get("device_tracker.device") + async def test_dont_track_wired_clients( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_device_registry diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index b5546b72bdd..18007998eba 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -14,11 +14,13 @@ from homeassistant.components.unifi.const import ( CONF_ALLOW_UPTIME_SENSORS, CONF_TRACK_CLIENTS, CONF_TRACK_DEVICES, + DOMAIN as UNIFI_DOMAIN, ) from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE, EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er +from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.entity_registry import RegistryEntryDisabler import homeassistant.util.dt as dt_util @@ -183,6 +185,37 @@ async def test_bandwidth_sensors( assert hass.states.get("sensor.wired_client_rx") is None assert hass.states.get("sensor.wired_client_tx") is None + # Enable option + + options[CONF_ALLOW_BANDWIDTH_SENSORS] = True + hass.config_entries.async_update_entry(config_entry, options=options.copy()) + await hass.async_block_till_done() + + assert len(hass.states.async_all()) == 5 + assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4 + assert hass.states.get("sensor.wireless_client_rx") + assert hass.states.get("sensor.wireless_client_tx") + assert hass.states.get("sensor.wired_client_rx") + assert hass.states.get("sensor.wired_client_tx") + + # Try to add the sensors again, using a signal + + clients_connected = {wired_client["mac"], wireless_client["mac"]} + devices_connected = set() + + controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + + async_dispatcher_send( + hass, + controller.signal_update, + clients_connected, + devices_connected, + ) + await hass.async_block_till_done() + + assert len(hass.states.async_all()) == 5 + assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4 + @pytest.mark.parametrize( ("initial_uptime", "event_uptime", "new_uptime"), @@ -267,6 +300,35 @@ async def test_uptime_sensors( assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 0 assert hass.states.get("sensor.client1_uptime") is None + # Enable option + + options[CONF_ALLOW_UPTIME_SENSORS] = True + with patch("homeassistant.util.dt.now", return_value=now): + hass.config_entries.async_update_entry(config_entry, options=options.copy()) + await hass.async_block_till_done() + + assert len(hass.states.async_all()) == 2 + assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1 + assert hass.states.get("sensor.client1_uptime") + + # Try to add the sensors again, using a signal + + clients_connected = {uptime_client["mac"]} + devices_connected = set() + + controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] + + async_dispatcher_send( + hass, + controller.signal_update, + clients_connected, + devices_connected, + ) + await hass.async_block_till_done() + + assert len(hass.states.async_all()) == 2 + assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 1 + async def test_remove_sensors( hass: HomeAssistant,