Registries store directly in data on loading.

homeassistant/helpers/area_registry.py

@@ -1,6 +1,7 @@
 """Provide a way to connect devices to one physical location."""
 import logging
 import uuid
+from asyncio import Event
 from collections import OrderedDict
 from typing import MutableMapping  # noqa: F401
 from typing import Iterable, Optional, cast
@@ -9,6 +10,7 @@
 
 from homeassistant.core import callback
 from homeassistant.loader import bind_hass
+
 from .typing import HomeAssistantType
 
 _LOGGER = logging.getLogger(__name__)
@@ -133,14 +135,23 @@ def _data_to_save(self) -> dict:
 @bind_hass
 async def async_get_registry(hass: HomeAssistantType) -> AreaRegistry:
     """Return area registry instance."""
-    task = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
+
+    if not reg_or_evt:
+        evt = hass.data[DATA_REGISTRY] = Event()
+
+        reg = AreaRegistry(hass)
+        await reg.async_load()
 
-    if task is None:
-        async def _load_reg() -> AreaRegistry:
-            registry = AreaRegistry(hass)
-            await registry.async_load()
-            return registry
+        hass.data[DATA_REGISTRY] = reg
+        evt.set()
+        return reg
 
-        task = hass.data[DATA_REGISTRY] = hass.async_create_task(_load_reg())
+    if isinstance(reg_or_evt, Event):
+        # This shouldn't be possible in production (has been tested)
+        # It's here soley for testing purposes.
+        evt = reg_or_evt
+        await evt.wait()
+        return cast(AreaRegistry, hass.data.get(DATA_REGISTRY))
 
-    return cast(AreaRegistry, await task)
+    return cast(AreaRegistry, reg_or_evt)

homeassistant/helpers/device_registry.py

@@ -1,15 +1,17 @@
 """Provide a way to connect entities belonging to one device."""
 import logging
 import uuid
-from typing import List, Optional
-
+from asyncio import Event
 from collections import OrderedDict
+from typing import List, Optional, cast
 
 import attr
 
 from homeassistant.core import callback
 from homeassistant.loader import bind_hass
 
+from .typing import HomeAssistantType
+
 _LOGGER = logging.getLogger(__name__)
 _UNDEF = object()
 
@@ -273,19 +275,28 @@ def async_clear_area_id(self, area_id: str) -> None:
 
 
 @bind_hass
-async def async_get_registry(hass) -> DeviceRegistry:
+async def async_get_registry(hass: HomeAssistantType) -> DeviceRegistry:
     """Return device registry instance."""
-    task = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
+
+    if not reg_or_evt:
+        evt = hass.data[DATA_REGISTRY] = Event()
+
+        reg = DeviceRegistry(hass)
+        await reg.async_load()
 
-    if task is None:
-        async def _load_reg():
-            registry = DeviceRegistry(hass)
-            await registry.async_load()
-            return registry
+        hass.data[DATA_REGISTRY] = reg
+        evt.set()
+        return reg
 
-        task = hass.data[DATA_REGISTRY] = hass.async_create_task(_load_reg())
+    if isinstance(reg_or_evt, Event):
+        # This shouldn't be possible in production (has been tested)
+        # It's here soley for testing purposes.
+        evt = reg_or_evt
+        await evt.wait()
+        return cast(DeviceRegistry, hass.data.get(DATA_REGISTRY))
 
-    return await task
+    return cast(DeviceRegistry, reg_or_evt)
 
 
 @callback

homeassistant/helpers/entity_registry.py

@@ -7,10 +7,11 @@
 registered. Registering a new entity while a timer is in progress resets the
 timer.
 """
+from asyncio import Event
 from collections import OrderedDict
 from itertools import chain
 import logging
-from typing import Optional, List
+from typing import List, Optional, cast
 import weakref
 
 import attr
@@ -20,6 +21,8 @@
 from homeassistant.util import ensure_unique_string, slugify
 from homeassistant.util.yaml import load_yaml
 
+from .typing import HomeAssistantType
+
 PATH_REGISTRY = 'entity_registry.yaml'
 DATA_REGISTRY = 'entity_registry'
 SAVE_DELAY = 10
@@ -277,19 +280,28 @@ def async_clear_config_entry(self, config_entry):
 
 
 @bind_hass
-async def async_get_registry(hass) -> EntityRegistry:
+async def async_get_registry(hass: HomeAssistantType) -> EntityRegistry:
     """Return entity registry instance."""
-    task = hass.data.get(DATA_REGISTRY)
+    reg_or_evt = hass.data.get(DATA_REGISTRY)
+
+    if not reg_or_evt:
+        evt = hass.data[DATA_REGISTRY] = Event()
+
+        reg = EntityRegistry(hass)
+        await reg.async_load()
 
-    if task is None:
-        async def _load_reg():
-            registry = EntityRegistry(hass)
-            await registry.async_load()
-            return registry
+        hass.data[DATA_REGISTRY] = reg
+        evt.set()
+        return reg
 
-        task = hass.data[DATA_REGISTRY] = hass.async_create_task(_load_reg())
+    if isinstance(reg_or_evt, Event):
+        # This shouldn't be possible in production (has been tested)
+        # It's here soley for testing purposes.
+        evt = reg_or_evt
+        await evt.wait()
+        return cast(EntityRegistry, hass.data.get(DATA_REGISTRY))
 
-    return await task
+    return cast(EntityRegistry, reg_or_evt)
 
 
 @callback

tests/common.py

@@ -327,11 +327,7 @@ def mock_registry(hass, mock_entries=None):
     registry = entity_registry.EntityRegistry(hass)
     registry.entities = mock_entries or OrderedDict()
 
-    async def _get_reg():
-        return registry
-
-    hass.data[entity_registry.DATA_REGISTRY] = \
-        hass.loop.create_task(_get_reg())
+    hass.data[entity_registry.DATA_REGISTRY] = registry
     return registry
 
 
@@ -340,11 +336,7 @@ def mock_area_registry(hass, mock_entries=None):
     registry = area_registry.AreaRegistry(hass)
     registry.areas = mock_entries or OrderedDict()
 
-    async def _get_reg():
-        return registry
-
-    hass.data[area_registry.DATA_REGISTRY] = \
-        hass.loop.create_task(_get_reg())
+    hass.data[area_registry.DATA_REGISTRY] = registry
     return registry
 
 
@@ -353,11 +345,7 @@ def mock_device_registry(hass, mock_entries=None):
     registry = device_registry.DeviceRegistry(hass)
     registry.devices = mock_entries or OrderedDict()
 
-    async def _get_reg():
-        return registry
-
-    hass.data[device_registry.DATA_REGISTRY] = \
-        hass.loop.create_task(_get_reg())
+    hass.data[device_registry.DATA_REGISTRY] = registry
     return registry