From 3ec7258c216796c70fb5e997be9492ed14f9e324 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 7 Nov 2020 06:54:20 -1000
Subject: [PATCH] Ensure instances of async_add_entities can be garbage
 collected (#42919)

---
 homeassistant/helpers/entity_platform.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/homeassistant/helpers/entity_platform.py b/homeassistant/helpers/entity_platform.py
index 693c01b982f..ddd1847f6a8 100644
--- a/homeassistant/helpers/entity_platform.py
+++ b/homeassistant/helpers/entity_platform.py
@@ -63,6 +63,8 @@ class EntityPlatform:
         self.config_entry: Optional[config_entries.ConfigEntry] = None
         self.entities: Dict[str, Entity] = {}  # pylint: disable=used-before-assignment
         self._tasks: List[asyncio.Future] = []
+        # Stop tracking tasks after setup is completed
+        self._setup_complete = False
         # Method to cancel the state change listener
         self._async_unsub_polling: Optional[CALLBACK_TYPE] = None
         # Method to cancel the retry of setup
@@ -197,7 +199,7 @@ class EntityPlatform:
                 await asyncio.shield(task)
 
             # Block till all entities are done
-            if self._tasks:
+            while self._tasks:
                 pending = [task for task in self._tasks if not task.done()]
                 self._tasks.clear()
 
@@ -205,6 +207,7 @@ class EntityPlatform:
                     await asyncio.gather(*pending)
 
             hass.config.components.add(full_name)
+            self._setup_complete = True
             return True
         except PlatformNotReady:
             tries += 1
@@ -258,14 +261,13 @@ class EntityPlatform:
         self, new_entities: Iterable["Entity"], update_before_add: bool = False
     ) -> None:
         """Schedule adding entities for a single platform async."""
-        self._tasks.append(
-            self.hass.async_create_task(
-                self.async_add_entities(
-                    new_entities, update_before_add=update_before_add
-                ),
-            )
+        task = self.hass.async_create_task(
+            self.async_add_entities(new_entities, update_before_add=update_before_add),
         )
 
+        if not self._setup_complete:
+            self._tasks.append(task)
+
     def add_entities(
         self, new_entities: Iterable["Entity"], update_before_add: bool = False
     ) -> None:
@@ -523,6 +525,7 @@ class EntityPlatform:
         if self._async_unsub_polling is not None:
             self._async_unsub_polling()
             self._async_unsub_polling = None
+        self._setup_complete = False
 
     async def async_destroy(self) -> None:
         """Destroy an entity platform.