From 3ed7a7166d65c9ed39cb35c61d18354cbe96d047 Mon Sep 17 00:00:00 2001 From: Jan Bouwhuis Date: Thu, 14 Mar 2024 22:11:44 +0100 Subject: [PATCH] Revert "Remove unused test helper mock_area_registry" (#113453) --- tests/common.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) 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,