Wait for device registry in entity registry loading (#166636)

homeassistant/bootstrap.py

@@ -468,6 +468,7 @@ async def async_load_base_functionality(hass: core.HomeAssistant) -> bool:
     translation.async_setup(hass)
 
     recovery = hass.config.recovery_mode
+    device_registry.async_setup(hass)
     try:
         await asyncio.gather(
             create_eager_task(get_internal_store_manager(hass).async_initialize()),

homeassistant/helpers/device_registry.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections import defaultdict
 from collections.abc import Iterable, Mapping
 from datetime import datetime
@@ -771,6 +772,7 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
     devices: ActiveDeviceRegistryItems
     deleted_devices: DeviceRegistryItems[DeletedDeviceEntry]
     _device_data: dict[str, DeviceEntry]
+    _loaded_event: asyncio.Event | None = None
 
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the device registry."""
@@ -784,6 +786,11 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
             serialize_in_event_loop=False,
         )
 
+    @callback
+    def async_setup(self) -> None:
+        """Set up the registry."""
+        self._loaded_event = asyncio.Event()
+
     @callback
     def async_get(self, device_id: str) -> DeviceEntry | None:
         """Get device.
@@ -1463,6 +1470,9 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
 
     async def _async_load(self) -> None:
         """Load the device registry."""
+        assert self._loaded_event is not None
+        assert not self._loaded_event.is_set()
+
         async_setup_cleanup(self.hass, self)
 
         data = await self._store.async_load()
@@ -1560,6 +1570,16 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
         self.deleted_devices = deleted_devices
         self._device_data = devices.data
 
+        self._loaded_event.set()
+
+    async def async_wait_loaded(self) -> None:
+        """Wait until the device registry is fully loaded.
+
+        Will only wait if the registry had already been set up.
+        """
+        if self._loaded_event is not None:
+            await self._loaded_event.wait()
+
     @callback
     def _data_to_save(self) -> dict[str, Any]:
         """Return data of device registry to store in a file."""
@@ -1706,9 +1726,14 @@ def async_get(hass: HomeAssistant) -> DeviceRegistry:
     return DeviceRegistry(hass)
 
 
+def async_setup(hass: HomeAssistant) -> None:
+    """Set up device registry."""
+    assert DATA_REGISTRY not in hass.data
+    async_get(hass).async_setup()
+
+
 async def async_load(hass: HomeAssistant, *, load_empty: bool = False) -> None:
     """Load device registry."""
-    assert DATA_REGISTRY not in hass.data
     await async_get(hass).async_load(load_empty=load_empty)
 
 

homeassistant/helpers/entity_registry.py

@@ -1944,6 +1944,10 @@ class EntityRegistry(BaseRegistry):
 
     async def _async_load(self) -> None:
         """Load the entity registry."""
+        # Device registry must be loaded before entity registry because
+        # migration and entity processing reference device names.
+        await dr.async_get(self.hass).async_wait_loaded()
+
         _async_setup_cleanup(self.hass, self)
         _async_setup_entity_restore(self.hass, self)
 

homeassistant/scripts/auth.py

@@ -55,6 +55,7 @@ def run(args: Sequence[str] | None) -> None:
 async def run_command(args: argparse.Namespace) -> None:
     """Run the command."""
     hass = HomeAssistant(os.path.join(os.getcwd(), args.config))
+    dr.async_setup(hass)
     await asyncio.gather(dr.async_load(hass), er.async_load(hass))
     hass.auth = await auth_manager_from_config(hass, [{"type": "homeassistant"}], [])
     provider = hass.auth.auth_providers[0]

homeassistant/scripts/check_config.py

@@ -302,6 +302,7 @@ async def async_check_config(config_dir):
     hass = core.HomeAssistant(config_dir)
     loader.async_setup(hass)
     hass.config_entries = ConfigEntries(hass, {})
+    dr.async_setup(hass)
     await ar.async_load(hass)
     await dr.async_load(hass)
     await er.async_load(hass)

tests/common.py

@@ -305,6 +305,8 @@ async def async_test_home_assistant(
         hass
     )
     if load_registries:
+        dr.async_setup(hass)
+
         with (
             patch.object(StoreWithoutWriteLoad, "async_load", return_value=None),
             patch(

tests/helpers/test_device_registry.py

@@ -363,6 +363,7 @@ async def test_loading_from_storage(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
     assert len(registry.devices) == 1
@@ -500,6 +501,7 @@ async def test_migration_from_1_1(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -654,6 +656,7 @@ async def test_migration_from_1_2(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -790,6 +793,7 @@ async def test_migration_fom_1_3(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -928,6 +932,7 @@ async def test_migration_from_1_4(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -1068,6 +1073,7 @@ async def test_migration_from_1_5(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -1210,6 +1216,7 @@ async def test_migration_from_1_6(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -1354,6 +1361,7 @@ async def test_migration_from_1_7(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -1496,6 +1504,7 @@ async def test_migration_from_1_10(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -1632,6 +1641,7 @@ async def test_migration_from_1_11(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
 
@@ -2627,6 +2637,7 @@ async def test_loading_saving_data(
     # Now load written data in new registry
     registry2 = dr.DeviceRegistry(hass)
     await flush_store(device_registry._store)
+    registry2.async_setup()
     await registry2.async_load()
 
     # Ensure same order
@@ -3782,6 +3793,7 @@ async def test_cleanup_entity_registry_change(
     Don't pre-load the registries as the debouncer will then not be waiting for
     EVENT_ENTITY_REGISTRY_UPDATED events.
     """
+    dr.async_setup(hass)
     await dr.async_load(hass)
     await er.async_load(hass)
     dev_reg = dr.async_get(hass)
@@ -4943,6 +4955,7 @@ async def test_loading_invalid_configuration_url_from_storage(
         },
     }
 
+    dr.async_setup(hass)
     await dr.async_load(hass)
     registry = dr.async_get(hass)
     assert len(registry.devices) == 1

tests/helpers/test_entity_registry.py

@@ -1,5 +1,6 @@
 """Tests for the Entity Registry."""
 
+import asyncio
 from datetime import datetime, timedelta
 from functools import partial
 from typing import Any
@@ -504,6 +505,49 @@ async def test_loading_saving_data(
     assert new_entry2.unit_of_measurement == "initial-unit_of_measurement"
 
 
+@pytest.mark.parametrize("load_registries", [False])
+async def test_entity_registry_loading_waits_for_device_registry(
+    hass: HomeAssistant, hass_storage: dict[str, Any]
+) -> None:
+    """Test entity registry waits for device registry when loaded concurrently.
+
+    Both registries are loaded in parallel during bootstrap via asyncio.gather.
+    The entity registry accesses device registry during loading. This test delays
+    the device registry store load so entity registry attempts to load first.
+    """
+    hass_storage[er.STORAGE_KEY] = {
+        "version": 1,
+        "minor_version": 1,
+        "data": {
+            "entities": [
+                {
+                    "entity_id": "test.my_entity",
+                    "device_id": "some-device",
+                    "platform": "test_platform",
+                    "unique_id": "unique-1",
+                },
+            ]
+        },
+    }
+
+    original_load = dr.DeviceRegistryStore.async_load
+
+    async def delayed_load(self: dr.DeviceRegistryStore) -> Any:
+        await asyncio.sleep(0)
+        return await original_load(self)
+
+    dr.async_setup(hass)
+
+    with patch.object(dr.DeviceRegistryStore, "async_load", delayed_load):
+        await asyncio.gather(
+            er.async_load(hass),
+            dr.async_load(hass),
+        )
+
+    registry = er.async_get(hass)
+    assert registry.async_get("test.my_entity") is not None
+
+
 def test_get_available_entity_id_considers_registered_entities(
     entity_registry: er.EntityRegistry,
 ) -> None:
@@ -1547,6 +1591,7 @@ async def test_migration_1_20(
             "deleted_devices": [],
         },
     }
+    dr.async_setup(hass)
     await dr.async_load(hass)
 
     # Entity registry data at version 1.20