From 1f1c66b3c2d5c279f92e2bd14050f5983e357d9f Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 19 Feb 2024 13:58:44 -0600
Subject: [PATCH] Group adding unifi entities together to reduce number of
 tasks (#110965)

---
 homeassistant/components/unifi/controller.py | 55 +++++++++++++-------
 1 file changed, 35 insertions(+), 20 deletions(-)

diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py
index 5873fa92cf7b3b..5ef8963e2acfbe 100644
--- a/homeassistant/components/unifi/controller.py
+++ b/homeassistant/components/unifi/controller.py
@@ -2,7 +2,9 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Iterable
 from datetime import datetime, timedelta
+from functools import partial
 import ssl
 from types import MappingProxyType
 from typing import Any, Literal
@@ -185,6 +187,17 @@ def register_platform(
             entity_class, descriptions, async_add_entities
         )
 
+    @callback
+    def _async_should_add_entity(
+        self, description: UnifiEntityDescription, obj_id: str
+    ) -> bool:
+        """Check if entity should be added."""
+        return bool(
+            (description.key, obj_id) not in self.known_objects
+            and description.allowed_fn(self, obj_id)
+            and description.supported_fn(self, obj_id)
+        )
+
     @callback
     def register_platform_add_entities(
         self,
@@ -195,45 +208,47 @@ def register_platform_add_entities(
         """Subscribe to UniFi API handlers and create entities."""
 
         @callback
-        def async_load_entities(description: UnifiEntityDescription) -> None:
+        def async_load_entities(descriptions: Iterable[UnifiEntityDescription]) -> None:
             """Load and subscribe to UniFi endpoints."""
-            api_handler = description.api_handler_fn(self.api)
 
             @callback
-            def async_add_unifi_entity(obj_ids: list[str]) -> None:
+            def async_add_unifi_entities() -> None:
                 """Add UniFi entity."""
                 async_add_entities(
                     [
                         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)
+                        for description in descriptions
+                        for obj_id in description.api_handler_fn(self.api)
+                        if self._async_should_add_entity(description, obj_id)
                     ]
                 )
 
-            async_add_unifi_entity(list(api_handler))
+            async_add_unifi_entities()
 
             @callback
-            def async_create_entity(event: ItemEvent, obj_id: str) -> None:
+            def async_create_entity(
+                description: UnifiEntityDescription, event: ItemEvent, obj_id: str
+            ) -> None:
                 """Create new UniFi entity on event."""
-                async_add_unifi_entity([obj_id])
-
-            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))
+                if self._async_should_add_entity(description, obj_id):
+                    async_add_entities(
+                        [unifi_platform_entity(obj_id, self, description)]
+                    )
+
+            for description in descriptions:
+                description.api_handler_fn(self.api).subscribe(
+                    partial(async_create_entity, description), ItemEvent.ADDED
+                )
 
             self.config_entry.async_on_unload(
                 async_dispatcher_connect(
-                    self.hass, self.signal_options_update, async_options_updated
+                    self.hass,
+                    self.signal_options_update,
+                    async_add_unifi_entities,
                 )
             )
 
-        for description in descriptions:
-            async_load_entities(description)
+        async_load_entities(descriptions)
 
     @property
     def signal_reachable(self) -> str: