Adjust registry access in tests root (#88880)

This commit is contained in:
epenet 2023-03-01 03:56:46 +01:00 committed by GitHub
parent 0e4c32efe2
commit 1bed5c7775
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -50,13 +50,13 @@ from homeassistant.core import (
callback,
)
from homeassistant.helpers import (
area_registry,
device_registry,
area_registry as ar,
device_registry as dr,
entity,
entity_platform,
entity_registry,
entity_registry as er,
intent,
issue_registry,
issue_registry as ir,
recorder as recorder_helper,
restore_state,
storage,
@ -251,10 +251,10 @@ async def async_test_home_assistant(event_loop, load_registries=True):
if load_registries:
with patch("homeassistant.helpers.storage.Store.async_load", return_value=None):
await asyncio.gather(
area_registry.async_load(hass),
device_registry.async_load(hass),
entity_registry.async_load(hass),
issue_registry.async_load(hass),
ar.async_load(hass),
dr.async_load(hass),
er.async_load(hass),
ir.async_load(hass),
)
hass.data[bootstrap.DATA_REGISTRIES_LOADED] = None
@ -481,8 +481,8 @@ def mock_component(hass: HomeAssistant, component: str) -> None:
def mock_registry(
hass: HomeAssistant,
mock_entries: dict[str, entity_registry.RegistryEntry] | None = None,
) -> entity_registry.EntityRegistry:
mock_entries: dict[str, er.RegistryEntry] | None = None,
) -> er.EntityRegistry:
"""Mock the Entity Registry.
This should only be used if you need to mock/re-stage a clean mocked
@ -494,20 +494,20 @@ def mock_registry(
If you just need to access the existing registry, use the `entity_registry`
fixture instead.
"""
registry = entity_registry.EntityRegistry(hass)
registry = er.EntityRegistry(hass)
if mock_entries is None:
mock_entries = {}
registry.entities = entity_registry.EntityRegistryItems()
registry.entities = er.EntityRegistryItems()
for key, entry in mock_entries.items():
registry.entities[key] = entry
hass.data[entity_registry.DATA_REGISTRY] = registry
hass.data[er.DATA_REGISTRY] = registry
return registry
def mock_area_registry(
hass: HomeAssistant, mock_entries: dict[str, area_registry.AreaEntry] | None = None
) -> area_registry.AreaRegistry:
hass: HomeAssistant, mock_entries: dict[str, ar.AreaEntry] | None = None
) -> ar.AreaRegistry:
"""Mock the Area Registry.
This should only be used if you need to mock/re-stage a clean mocked
@ -519,17 +519,17 @@ def mock_area_registry(
If you just need to access the existing registry, use the `area_registry`
fixture instead.
"""
registry = area_registry.AreaRegistry(hass)
registry = ar.AreaRegistry(hass)
registry.areas = mock_entries or OrderedDict()
hass.data[area_registry.DATA_REGISTRY] = registry
hass.data[ar.DATA_REGISTRY] = registry
return registry
def mock_device_registry(
hass: HomeAssistant,
mock_entries: dict[str, device_registry.DeviceEntry] | None = None,
) -> device_registry.DeviceRegistry:
mock_entries: dict[str, dr.DeviceEntry] | None = None,
) -> dr.DeviceRegistry:
"""Mock the Device Registry.
This should only be used if you need to mock/re-stage a clean mocked
@ -541,15 +541,15 @@ def mock_device_registry(
If you just need to access the existing registry, use the `device_registry`
fixture instead.
"""
registry = device_registry.DeviceRegistry(hass)
registry.devices = device_registry.DeviceRegistryItems()
registry = dr.DeviceRegistry(hass)
registry.devices = dr.DeviceRegistryItems()
if mock_entries is None:
mock_entries = {}
for key, entry in mock_entries.items():
registry.devices[key] = entry
registry.deleted_devices = device_registry.DeviceRegistryItems()
registry.deleted_devices = dr.DeviceRegistryItems()
hass.data[device_registry.DATA_REGISTRY] = registry
hass.data[dr.DATA_REGISTRY] = registry
return registry