Implement ConfigEntry async_wait_for_states (#81771)

* Implement async_wait_for_states

* stale docstr, remove hints

* Assert return states for tests

tests/test_config_entries.py

@@ -3330,3 +3330,110 @@ async def test_reauth(hass):
     entry2.async_start_reauth(hass, {"extra_context": "some_extra_context"})
     await hass.async_block_till_done()
     assert len(hass.config_entries.flow.async_progress()) == 2
+
+
+async def test_wait_for_loading_entry(hass):
+    """Test waiting for entry to be set up."""
+
+    entry = MockConfigEntry(title="test_title", domain="test")
+
+    mock_setup_entry = AsyncMock(return_value=True)
+    mock_integration(hass, MockModule("test", async_setup_entry=mock_setup_entry))
+    mock_entity_platform(hass, "config_flow.test", None)
+
+    await entry.async_setup(hass)
+    await hass.async_block_till_done()
+
+    flow = hass.config_entries.flow
+
+    async def _load_entry():
+        # Mock config entry
+        assert await async_setup_component(hass, "test", {})
+
+    entry = MockConfigEntry(title="test_title", domain="test")
+    entry.add_to_hass(hass)
+    flow = hass.config_entries.flow
+    with patch.object(flow, "async_init", wraps=flow.async_init):
+        hass.async_add_job(_load_entry)
+        new_state = await hass.config_entries.async_wait_for_states(
+            entry,
+            {
+                config_entries.ConfigEntryState.LOADED,
+                config_entries.ConfigEntryState.SETUP_ERROR,
+            },
+            timeout=1.0,
+        )
+        assert new_state is config_entries.ConfigEntryState.LOADED
+    assert entry.state is config_entries.ConfigEntryState.LOADED
+
+
+async def test_wait_for_loading_failed_entry(hass):
+    """Test waiting for entry to be set up that fails loading."""
+
+    entry = MockConfigEntry(title="test_title", domain="test")
+
+    mock_setup_entry = AsyncMock(side_effect=HomeAssistantError)
+    mock_integration(hass, MockModule("test", async_setup_entry=mock_setup_entry))
+    mock_entity_platform(hass, "config_flow.test", None)
+
+    await entry.async_setup(hass)
+    await hass.async_block_till_done()
+
+    flow = hass.config_entries.flow
+
+    async def _load_entry():
+        # Mock config entry
+        assert await async_setup_component(hass, "test", {})
+
+    entry = MockConfigEntry(title="test_title", domain="test")
+    entry.add_to_hass(hass)
+    flow = hass.config_entries.flow
+    with patch.object(flow, "async_init", wraps=flow.async_init):
+        hass.async_add_job(_load_entry)
+        new_state = await hass.config_entries.async_wait_for_states(
+            entry,
+            {
+                config_entries.ConfigEntryState.LOADED,
+                config_entries.ConfigEntryState.SETUP_ERROR,
+            },
+            timeout=1.0,
+        )
+        assert new_state is config_entries.ConfigEntryState.SETUP_ERROR
+    assert entry.state is config_entries.ConfigEntryState.SETUP_ERROR
+
+
+async def test_wait_for_loading_timeout(hass):
+    """Test waiting for entry to be set up that fails with a timeout."""
+
+    async def _async_setup_entry(hass, entry):
+        await asyncio.sleep(1)
+        return True
+
+    entry = MockConfigEntry(title="test_title", domain="test")
+
+    mock_integration(hass, MockModule("test", async_setup_entry=_async_setup_entry))
+    mock_entity_platform(hass, "config_flow.test", None)
+
+    await entry.async_setup(hass)
+    await hass.async_block_till_done()
+
+    flow = hass.config_entries.flow
+
+    async def _load_entry():
+        # Mock config entry
+        assert await async_setup_component(hass, "test", {})
+
+    entry = MockConfigEntry(title="test_title", domain="test")
+    entry.add_to_hass(hass)
+    flow = hass.config_entries.flow
+    with patch.object(flow, "async_init", wraps=flow.async_init):
+        hass.async_add_job(_load_entry)
+        with pytest.raises(asyncio.exceptions.TimeoutError):
+            await hass.config_entries.async_wait_for_states(
+                entry,
+                {
+                    config_entries.ConfigEntryState.LOADED,
+                    config_entries.ConfigEntryState.SETUP_ERROR,
+                },
+                timeout=0.1,
+            )