diff --git a/tests/common.py b/tests/common.py index 6d0e74021e3..834354fa673 100644 --- a/tests/common.py +++ b/tests/common.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from collections import OrderedDict from collections.abc import AsyncGenerator, Generator, Mapping, Sequence from contextlib import asynccontextmanager, contextmanager from datetime import UTC, datetime, timedelta @@ -620,6 +621,27 @@ def mock_registry( return registry +def mock_area_registry( + 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 + area registry in your current hass object. It can be useful to, + for example, pre-load the registry with items. + + This mock will thus replace the existing registry in the running hass. + + If you just need to access the existing registry, use the `area_registry` + fixture instead. + """ + registry = ar.AreaRegistry(hass) + registry.areas = mock_entries or OrderedDict() + + hass.data[ar.DATA_REGISTRY] = registry + return registry + + def mock_device_registry( hass: HomeAssistant, mock_entries: dict[str, dr.DeviceEntry] | None = None,