Allow loading UniFi entities on config options change (#88762)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>

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)
 

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(

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

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,