Use eager task creation to add entities to entity platform (#111553)

This commit is contained in:
J. Nick Koston 2024-02-27 09:54:51 -10:00 committed by GitHub
parent 8da2c53742
commit a6f4f6eae8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 10 additions and 4 deletions

View File

@ -476,6 +476,7 @@ class EntityPlatform:
task = self.hass.async_create_task( task = self.hass.async_create_task(
self.async_add_entities(new_entities, update_before_add=update_before_add), self.async_add_entities(new_entities, update_before_add=update_before_add),
f"EntityPlatform async_add_entities {self.domain}.{self.platform_name}", f"EntityPlatform async_add_entities {self.domain}.{self.platform_name}",
eager_start=True,
) )
if not self._setup_complete: if not self._setup_complete:
@ -491,6 +492,7 @@ class EntityPlatform:
self.hass, self.hass,
self.async_add_entities(new_entities, update_before_add=update_before_add), self.async_add_entities(new_entities, update_before_add=update_before_add),
f"EntityPlatform async_add_entities_for_entry {self.domain}.{self.platform_name}", f"EntityPlatform async_add_entities_for_entry {self.domain}.{self.platform_name}",
eager_start=True,
) )
if not self._setup_complete: if not self._setup_complete:
@ -526,9 +528,10 @@ class EntityPlatform:
event loop and will finish faster if we run them concurrently. event loop and will finish faster if we run them concurrently.
""" """
results: list[BaseException | None] | None = None results: list[BaseException | None] | None = None
tasks = [create_eager_task(coro) for coro in coros]
try: try:
async with self.hass.timeout.async_timeout(timeout, self.domain): async with self.hass.timeout.async_timeout(timeout, self.domain):
results = await asyncio.gather(*coros, return_exceptions=True) results = await asyncio.gather(*tasks, return_exceptions=True)
except TimeoutError: except TimeoutError:
self.logger.warning( self.logger.warning(
"Timed out adding entities for domain %s with platform %s after %ds", "Timed out adding entities for domain %s with platform %s after %ds",

View File

@ -195,15 +195,18 @@ async def test_remove_stale_device(
(DOMAIN, 7654321) in device.identifiers for device in device_entries_other (DOMAIN, 7654321) in device.identifiers for device in device_entries_other
) )
assert await config_entry.async_unload(hass) assert await hass.config_entries.async_unload(config_entry.entry_id)
await hass.async_block_till_done() await hass.async_block_till_done()
assert config_entry.state == ConfigEntryState.NOT_LOADED assert config_entry.state is ConfigEntryState.NOT_LOADED
del location.devices_by_id[another_device.deviceid] del location.devices_by_id[another_device.deviceid]
await hass.config_entries.async_setup(config_entry.entry_id) await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done() await hass.async_block_till_done()
await hass.async_block_till_done()
assert config_entry.state is ConfigEntryState.LOADED assert config_entry.state is ConfigEntryState.LOADED
assert ( assert (
hass.states.async_entity_ids_count() == 3 hass.states.async_entity_ids_count() == 3
) # 1 climate entities; 2 sensor entities ) # 1 climate entities; 2 sensor entities

View File

@ -121,7 +121,7 @@ async def test_setup_component_with_config(
await hass.async_block_till_done() await hass.async_block_till_done()
assert fake_post_hits == 10 assert fake_post_hits >= 8
mock_impl.assert_called_once() mock_impl.assert_called_once()
mock_webhook.assert_called_once() mock_webhook.assert_called_once()