V2 timeout for async_add_entities (#38601)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
Co-authored-by: J. Nick Koston <nick@koston.org>
Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/helpers/entity_platform.py

@@ -23,6 +23,9 @@ if TYPE_CHECKING:
 
 SLOW_SETUP_WARNING = 10
 SLOW_SETUP_MAX_WAIT = 60
+SLOW_ADD_ENTITY_MAX_WAIT = 10  # Per Entity
+SLOW_ADD_MIN_TIMEOUT = 60
+
 PLATFORM_NOT_READY_RETRIES = 10
 DATA_ENTITY_PLATFORM = "entity_platform"
 PLATFORM_NOT_READY_BASE_WAIT_TIME = 30  # seconds
@@ -292,7 +295,17 @@ class EntityPlatform:
         if not tasks:
             return
 
-        await asyncio.gather(*tasks)
+        timeout = max(SLOW_ADD_ENTITY_MAX_WAIT * len(tasks), SLOW_ADD_MIN_TIMEOUT)
+        try:
+            async with self.hass.timeout.async_timeout(timeout, self.domain):
+                await asyncio.gather(*tasks)
+        except asyncio.TimeoutError:
+            self.logger.warning(
+                "Timed out adding entities for domain %s with platform %s after %ds",
+                self.domain,
+                self.platform_name,
+                timeout,
+            )
 
         if self._async_unsub_polling is not None or not any(
             entity.should_poll for entity in self.entities.values()

tests/helpers/test_entity_platform.py

@@ -931,3 +931,41 @@ async def test_invalid_entity_id(hass):
         await platform.async_add_entities([entity])
     assert entity.hass is None
     assert entity.platform is None
+
+
+class MockBlockingEntity(MockEntity):
+    """Class to mock an entity that will block adding entities."""
+
+    async def async_added_to_hass(self):
+        """Block for a long time."""
+        await asyncio.sleep(1000)
+
+
+async def test_setup_entry_with_entities_that_block_forever(hass, caplog):
+    """Test we cancel adding entities when we reach the timeout."""
+    registry = mock_registry(hass)
+
+    async def async_setup_entry(hass, config_entry, async_add_entities):
+        """Mock setup entry method."""
+        async_add_entities([MockBlockingEntity(name="test1", unique_id="unique")])
+        return True
+
+    platform = MockPlatform(async_setup_entry=async_setup_entry)
+    config_entry = MockConfigEntry(entry_id="super-mock-id")
+    mock_entity_platform = MockEntityPlatform(
+        hass, platform_name=config_entry.domain, platform=platform
+    )
+
+    with patch.object(entity_platform, "SLOW_ADD_ENTITY_MAX_WAIT", 0.01), patch.object(
+        entity_platform, "SLOW_ADD_MIN_TIMEOUT", 0.01
+    ):
+        assert await mock_entity_platform.async_setup_entry(config_entry)
+        await hass.async_block_till_done()
+    full_name = f"{mock_entity_platform.domain}.{config_entry.domain}"
+    assert full_name in hass.config.components
+    assert len(hass.states.async_entity_ids()) == 0
+    assert len(registry.entities) == 1
+    assert "Timed out adding entities" in caplog.text
+    assert "test_domain.test1" in caplog.text
+    assert "test_domain" in caplog.text
+    assert "test" in caplog.text