From 9a342f87c0b78ceccc733e27fe53e7f2ef6e20a4 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 7 Apr 2024 15:25:55 -1000
Subject: [PATCH] Avoid checking for polling if an entity fails to add
 (#115159)

* Avoid checking for polling if an entity fails to add

* no need to do protected access

* no need to do protected access

* no need to do protected access

* no need to do protected access

* coverage

* fix test

* fix

* broken one must be first
---
 homeassistant/helpers/entity_platform.py | 11 +++++++-
 tests/helpers/test_entity_platform.py    | 36 +++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/homeassistant/helpers/entity_platform.py b/homeassistant/helpers/entity_platform.py
index 1cff472af72e2f..6d7ed7ed1b8597 100644
--- a/homeassistant/helpers/entity_platform.py
+++ b/homeassistant/helpers/entity_platform.py
@@ -631,7 +631,16 @@ async def async_add_entities(
         if (
             (self.config_entry and self.config_entry.pref_disable_polling)
             or self._async_unsub_polling is not None
-            or not any(entity.should_poll for entity in entities)
+            or not any(
+                # Entity may have failed to add or called `add_to_platform_abort`
+                # so we check if the entity is in self.entities before
+                # checking `entity.should_poll` since `should_poll` may need to
+                # check `self.hass` which will be `None` if the entity did not add
+                entity.entity_id
+                and entity.entity_id in self.entities
+                and entity.should_poll
+                for entity in entities
+            )
         ):
             return
 
diff --git a/tests/helpers/test_entity_platform.py b/tests/helpers/test_entity_platform.py
index 31c6f8e6e30c49..59c4f7357f3df5 100644
--- a/tests/helpers/test_entity_platform.py
+++ b/tests/helpers/test_entity_platform.py
@@ -5,7 +5,7 @@
 from datetime import timedelta
 import logging
 from typing import Any
-from unittest.mock import ANY, Mock, patch
+from unittest.mock import ANY, AsyncMock, Mock, patch
 
 import pytest
 
@@ -78,6 +78,40 @@ async def test_polling_only_updates_entities_it_should_poll(
     assert poll_ent.async_update.called
 
 
+async def test_polling_check_works_if_entity_add_fails(
+    hass: HomeAssistant,
+) -> None:
+    """Test the polling check works if an entity add fails."""
+    component = EntityComponent(_LOGGER, DOMAIN, hass, timedelta(seconds=20))
+    await component.async_setup({})
+
+    class MockEntityNeedsSelfHassInShouldPoll(MockEntity):
+        """Mock entity that needs self.hass in should_poll."""
+
+        @property
+        def should_poll(self) -> bool:
+            """Return True if entity has to be polled."""
+            return self.hass.data is not None
+
+    working_poll_ent = MockEntityNeedsSelfHassInShouldPoll(should_poll=True)
+    working_poll_ent.async_update = AsyncMock()
+    broken_poll_ent = MockEntityNeedsSelfHassInShouldPoll(should_poll=True)
+    broken_poll_ent.async_update = AsyncMock(side_effect=Exception("Broken"))
+
+    await component.async_add_entities(
+        [broken_poll_ent, working_poll_ent], update_before_add=True
+    )
+
+    working_poll_ent.async_update.reset_mock()
+    broken_poll_ent.async_update.reset_mock()
+
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=20))
+    await hass.async_block_till_done(wait_background_tasks=True)
+
+    assert not broken_poll_ent.async_update.called
+    assert working_poll_ent.async_update.called
+
+
 async def test_polling_disabled_by_config_entry(hass: HomeAssistant) -> None:
     """Test the polling of only updated entities."""
     entity_platform = MockEntityPlatform(hass)