UniFi - Move some preloading of unavailable clients earlier in setup phase (#34599)

Improve readability of setting up switch platform

homeassistant/components/unifi/controller.py

@@ -265,6 +265,30 @@ class UniFiController:
             LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
             return False
 
+        # Restore clients that is not a part of active clients list.
+        entity_registry = await self.hass.helpers.entity_registry.async_get_registry()
+        for entity in entity_registry.entities.values():
+            if (
+                entity.config_entry_id != self.config_entry.entry_id
+                or "-" not in entity.unique_id
+            ):
+                continue
+
+            mac = ""
+            if entity.domain == TRACKER_DOMAIN:
+                mac, _ = entity.unique_id.split("-", 1)
+            elif entity.domain == SWITCH_DOMAIN:
+                _, mac = entity.unique_id.split("-", 1)
+
+            if mac in self.api.clients or mac not in self.api.clients_all:
+                continue
+
+            client = self.api.clients_all[mac]
+            self.api.clients.process_raw([client.raw])
+            LOGGER.debug(
+                "Restore disconnected client %s (%s)", entity.entity_id, client.mac,
+            )
+
         wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
         self.wireless_clients = wireless_clients.get_data(self.config_entry)
         self.update_wireless_clients()

homeassistant/components/unifi/device_tracker.py

@@ -47,26 +47,6 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
     controller.entities[DOMAIN] = {CLIENT_TRACKER: set(), DEVICE_TRACKER: set()}
 
-    # Restore clients that is not a part of active clients list.
-    entity_registry = await hass.helpers.entity_registry.async_get_registry()
-    for entity in entity_registry.entities.values():
-
-        if (
-            entity.config_entry_id == config_entry.entry_id
-            and entity.domain == DOMAIN
-            and "-" in entity.unique_id
-        ):
-
-            mac, _ = entity.unique_id.split("-", 1)
-            if mac in controller.api.clients or mac not in controller.api.clients_all:
-                continue
-
-            client = controller.api.clients_all[mac]
-            controller.api.clients.process_raw([client.raw])
-            LOGGER.debug(
-                "Restore disconnected client %s (%s)", entity.entity_id, client.mac,
-            )
-
     @callback
     def items_added():
         """Update the values of the controller."""

homeassistant/components/unifi/switch.py

@@ -30,63 +30,53 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     if controller.site_role != "admin":
         return
 
-    switches_off = []
-
-    # Restore clients that is not a part of active clients list.
+    # Store previously known POE control entities in case their POE are turned off.
+    previously_known_poe_clients = []
     entity_registry = await hass.helpers.entity_registry.async_get_registry()
     for entity in entity_registry.entities.values():
 
         if (
-            entity.config_entry_id == config_entry.entry_id
-            and entity.unique_id.startswith(f"{POE_SWITCH}-")
+            entity.config_entry_id != config_entry.entry_id
+            or not entity.unique_id.startswith(POE_SWITCH)
         ):
+            continue
 
-            _, mac = entity.unique_id.split("-", 1)
+        mac = entity.unique_id.replace(f"{POE_SWITCH}-", "")
+        if mac in controller.api.clients or mac in controller.api.clients_all:
+            previously_known_poe_clients.append(entity.unique_id)
 
-            if mac in controller.api.clients:
-                switches_off.append(entity.unique_id)
-                continue
-
-            if mac in controller.api.clients_all:
-                client = controller.api.clients_all[mac]
-                controller.api.clients.process_raw([client.raw])
-                switches_off.append(entity.unique_id)
-                continue
+    for mac in controller.option_block_clients:
+        if mac not in controller.api.clients and mac in controller.api.clients_all:
+            client = controller.api.clients_all[mac]
+            controller.api.clients.process_raw([client.raw])
 
     @callback
     def items_added():
         """Update the values of the controller."""
         if controller.option_block_clients or controller.option_poe_clients:
-            add_entities(controller, async_add_entities, switches_off)
+            add_entities(controller, async_add_entities, previously_known_poe_clients)
 
     for signal in (controller.signal_update, controller.signal_options_update):
         controller.listeners.append(async_dispatcher_connect(hass, signal, items_added))
 
     items_added()
-    switches_off.clear()
+    previously_known_poe_clients.clear()
 
 
 @callback
-def add_entities(controller, async_add_entities, switches_off):
+def add_entities(controller, async_add_entities, previously_known_poe_clients):
     """Add new switch entities from the controller."""
     switches = []
 
     for mac in controller.option_block_clients:
 
-        if mac in controller.entities[DOMAIN][BLOCK_SWITCH]:
-            continue
-
-        client = None
-
-        if mac in controller.api.clients:
-            client = controller.api.clients[mac]
-
-        elif mac in controller.api.clients_all:
-            client = controller.api.clients_all[mac]
-
-        if not client:
+        if (
+            mac in controller.entities[DOMAIN][BLOCK_SWITCH]
+            or mac not in controller.api.clients
+        ):
             continue
 
+        client = controller.api.clients[mac]
         switches.append(UniFiBlockClientSwitch(client, controller))
 
     if controller.option_poe_clients:
@@ -101,7 +91,7 @@ def add_entities(controller, async_add_entities, switches_off):
 
             client = controller.api.clients[mac]
 
-            if poe_client_id not in switches_off and (
+            if poe_client_id not in previously_known_poe_clients and (
                 mac in controller.wireless_clients
                 or client.sw_mac not in devices
                 or not devices[client.sw_mac].ports[client.sw_port].port_poe
@@ -114,7 +104,7 @@ def add_entities(controller, async_add_entities, switches_off):
             multi_clients_on_port = False
             for client2 in controller.api.clients.values():
 
-                if poe_client_id in switches_off:
+                if poe_client_id in previously_known_poe_clients:
                     break
 
                 if (

tests/components/unifi/test_switch.py

@@ -383,7 +383,7 @@ async def test_option_block_clients(hass):
         options={CONF_BLOCK_CLIENT: [BLOCKED["mac"], UNBLOCKED["mac"]]},
     )
     await hass.async_block_till_done()
-    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
 
     # Remove the second switch again
     hass.config_entries.async_update_entry(
@@ -397,7 +397,7 @@ async def test_option_block_clients(hass):
         controller.config_entry, options={CONF_BLOCK_CLIENT: [UNBLOCKED["mac"]]},
     )
     await hass.async_block_till_done()
-    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
+    assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
 
     # Remove one
     hass.config_entries.async_update_entry(