Merge branch 'dev' into add_restored_entities_to_entity_sources

homeassistant/helpers/entity.py

@@ -1436,6 +1436,10 @@ class Entity(
             # Set the entity's state will to unavailable + ATTR_RESTORED: True
             self.registry_entry.write_unavailable_state(self.hass)
         else:
+            # The check for self.platform guards against integrations not using an
+            # EntityComponent and can be removed in HA Core 2026.8
+            if self.platform:
+                del entity_sources(self.hass)[self.entity_id]
             self.hass.states.async_remove(self.entity_id, context=self._context)
 
     async def async_added_to_hass(self) -> None:
@@ -1495,10 +1499,6 @@ class Entity(
 
         Not to be extended by integrations.
         """
-        # The check for self.platform guards against integrations not using an
-        # EntityComponent and can be removed in HA Core 2026.8
-        if self.platform:
-            del entity_sources(self.hass)[self.entity_id]
 
     @callback
     def _async_registry_updated(

homeassistant/helpers/entity_registry.py

@@ -1896,6 +1896,9 @@ def _async_setup_cleanup(hass: HomeAssistant, registry: EntityRegistry) -> None:
 def _async_setup_entity_restore(hass: HomeAssistant, registry: EntityRegistry) -> None:
     """Set up the entity restore mechanism."""
 
+    # pylint: disable-next=import-outside-toplevel
+    from . import entity
+
     @callback
     def cleanup_restored_states_filter(event_data: Mapping[str, Any]) -> bool:
         """Clean up restored states filter."""
@@ -1923,6 +1926,7 @@ def _async_setup_entity_restore(hass: HomeAssistant, registry: EntityRegistry) -
         if state is None or not state.attributes.get(ATTR_RESTORED):
             return
 
+        del entity.entity_sources(hass)[event.data["entity_id"]]
         hass.states.async_remove(event.data["entity_id"], context=event.context)
 
     hass.bus.async_listen(
@@ -1939,10 +1943,18 @@ def _async_setup_entity_restore(hass: HomeAssistant, registry: EntityRegistry) -
         """Make sure state machine contains entry for each registered entity."""
         existing = set(hass.states.async_entity_ids())
 
+        entity_sources = entity.entity_sources(hass)
         for entry in registry.entities.values():
             if entry.entity_id in existing or entry.disabled:
                 continue
 
+            entity_info: entity.EntityInfo = {
+                "domain": entry.platform,
+            }
+            if entry.config_entry_id:
+                entity_info["config_entry"] = entry.config_entry_id
+
+            entity_sources[entry.entity_id] = entity_info
             entry.write_unavailable_state(hass)
 
     hass.bus.async_listen(EVENT_HOMEASSISTANT_START, _write_unavailable_states)

tests/helpers/test_entity.py

@@ -823,29 +823,68 @@ async def test_warn_slow_write_state_custom_component(
     ) in caplog.text
 
 
-async def test_setup_source(hass: HomeAssistant) -> None:
+async def test_setup_source(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
     """Check that we register sources correctly."""
     platform = MockEntityPlatform(hass)
 
-    entity_platform = MockEntity(name="Platform Config Source")
-    await platform.async_add_entities([entity_platform])
+    entity_platform = MockEntity(name="Platform Config Source 1")
+    entity_platform_uniq = MockEntity(name="Platform Config Source 2", unique_id="ABC")
+    await platform.async_add_entities([entity_platform, entity_platform_uniq])
 
     platform.config_entry = MockConfigEntry()
-    entity_entry = MockEntity(name="Config Entry Source")
-    await platform.async_add_entities([entity_entry])
+    platform.config_entry.add_to_hass(hass)
+    entity_entry = MockEntity(name="Config Entry Source 1")
+    entity_entry_uniq = MockEntity(name="Config Entry Source 2", unique_id="DEF")
+    await platform.async_add_entities([entity_entry, entity_entry_uniq])
 
+    assert {state.entity_id: state.state for state in hass.states.async_all()} == {
+        "test_domain.config_entry_source_1": "unknown",
+        "test_domain.config_entry_source_2": "unknown",
+        "test_domain.platform_config_source_1": "unknown",
+        "test_domain.platform_config_source_2": "unknown",
+    }
     assert entity.entity_sources(hass) == {
-        "test_domain.platform_config_source": {
+        "test_domain.platform_config_source_1": {
             "domain": "test_platform",
         },
-        "test_domain.config_entry_source": {
+        "test_domain.platform_config_source_2": {
+            "domain": "test_platform",
+        },
+        "test_domain.config_entry_source_1": {
+            "config_entry": platform.config_entry.entry_id,
+            "domain": "test_platform",
+        },
+        "test_domain.config_entry_source_2": {
             "config_entry": platform.config_entry.entry_id,
             "domain": "test_platform",
         },
     }
 
+    # Reset the platform to remove all the entities
     await platform.async_reset()
 
+    # Check that entities with unique IDs are still present
+    assert {state.entity_id: state.state for state in hass.states.async_all()} == {
+        "test_domain.config_entry_source_2": "unavailable",
+        "test_domain.platform_config_source_2": "unavailable",
+    }
+    assert entity.entity_sources(hass) == {
+        "test_domain.platform_config_source_2": {
+            "domain": "test_platform",
+        },
+        "test_domain.config_entry_source_2": {
+            "config_entry": platform.config_entry.entry_id,
+            "domain": "test_platform",
+        },
+    }
+
+    # Remove the entity registry entries for the entities with unique IDs
+    entity_registry.async_remove("test_domain.platform_config_source_2")
+    entity_registry.async_remove("test_domain.config_entry_source_2")
+
+    assert hass.states.async_all() == []
     assert entity.entity_sources(hass) == {}
 
 

tests/helpers/test_entity_registry.py

@@ -18,7 +18,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import CoreState, Event, HomeAssistant, callback
 from homeassistant.exceptions import MaxLengthExceeded
-from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity, entity_registry as er
 from homeassistant.helpers.event import async_track_entity_registry_updated_event
 from homeassistant.helpers.typing import UNDEFINED
 from homeassistant.util.dt import utc_from_timestamp, utcnow
@@ -1962,6 +1962,7 @@ async def test_restore_states(
     hass: HomeAssistant, entity_registry: er.EntityRegistry
 ) -> None:
     """Test restoring states."""
+    entity_sources = entity.entity_sources(hass)
     hass.set_state(CoreState.not_running)
 
     entity_registry.async_get_or_create(
@@ -1978,18 +1979,23 @@ async def test_restore_states(
         suggested_object_id="disabled",
         disabled_by=er.RegistryEntryDisabler.HASS,
     )
+    config_entry = MockConfigEntry(domain="hue")
+    config_entry.add_to_hass(hass)
     entity_registry.async_get_or_create(
         "light",
         "hue",
         "9012",
         suggested_object_id="all_info_set",
         capabilities={"max": 100},
+        config_entry=config_entry,
         supported_features=5,
         original_device_class="mock-device-class",
         original_name="Mock Original Name",
         original_icon="hass:original-icon",
     )
 
+    assert entity_sources == {}
+
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START, {})
     await hass.async_block_till_done()
 
@@ -2013,6 +2019,16 @@ async def test_restore_states(
         "icon": "hass:original-icon",
     }
 
+    assert entity_sources == {
+        "light.all_info_set": {
+            "config_entry": config_entry.entry_id,
+            "domain": "hue",
+        },
+        "light.simple": {
+            "domain": "hue",
+        },
+    }
+
     entity_registry.async_remove("light.disabled")
     entity_registry.async_remove("light.simple")
     entity_registry.async_remove("light.all_info_set")
@@ -2023,6 +2039,8 @@ async def test_restore_states(
     assert hass.states.get("light.disabled") is None
     assert hass.states.get("light.all_info_set") is None
 
+    assert entity_sources == {}
+
 
 async def test_remove_device_removes_entities(
     hass: HomeAssistant,