diff --git a/homeassistant/helpers/area_registry.py b/homeassistant/helpers/area_registry.py index aeb52e8faed..e8ba4ce362d 100644 --- a/homeassistant/helpers/area_registry.py +++ b/homeassistant/helpers/area_registry.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections import OrderedDict from collections.abc import Container, Iterable, MutableMapping -from typing import Optional, cast +from typing import Any, Optional, cast import attr @@ -19,7 +19,8 @@ from .typing import UNDEFINED, UndefinedType DATA_REGISTRY = "area_registry" EVENT_AREA_REGISTRY_UPDATED = "area_registry_updated" STORAGE_KEY = "core.area_registry" -STORAGE_VERSION = 1 +STORAGE_VERSION_MAJOR = 1 +STORAGE_VERSION_MINOR = 2 SAVE_DELAY = 10 @@ -42,6 +43,28 @@ class AreaEntry: object.__setattr__(self, "id", suggestion) +class AreaRegistryStore(Store[dict[str, list[dict[str, Optional[str]]]]]): + """Store area registry data.""" + + async def _async_migrate_func( + self, + old_major_version: int, + old_minor_version: int, + old_data: dict[str, list[dict[str, str | None]]], + ) -> dict[str, Any]: + """Migrate to the new version.""" + if old_major_version < 2: + if old_minor_version < 2: + # Version 1.2 implements migration and freezes the available keys + for area in old_data["areas"]: + # Populate keys which were introduced before version 1.2 + area.setdefault("picture", None) + + if old_major_version > 1: + raise NotImplementedError + return old_data + + class AreaRegistry: """Class to hold a registry of areas.""" @@ -49,8 +72,12 @@ class AreaRegistry: """Initialize the area registry.""" self.hass = hass self.areas: MutableMapping[str, AreaEntry] = {} - self._store = Store[dict[str, list[dict[str, Optional[str]]]]]( - hass, STORAGE_VERSION, STORAGE_KEY, atomic_writes=True + self._store = AreaRegistryStore( + hass, + STORAGE_VERSION_MAJOR, + STORAGE_KEY, + atomic_writes=True, + minor_version=STORAGE_VERSION_MINOR, ) self._normalized_name_area_idx: dict[str, str] = {} @@ -183,11 +210,10 @@ class AreaRegistry: assert area["name"] is not None and area["id"] is not None normalized_name = normalize_area_name(area["name"]) areas[area["id"]] = AreaEntry( - name=area["name"], id=area["id"], - # New in 2021.11 - picture=area.get("picture"), + name=area["name"], normalized_name=normalized_name, + picture=area["picture"], ) self._normalized_name_area_idx[normalized_name] = area["id"] diff --git a/tests/helpers/test_area_registry.py b/tests/helpers/test_area_registry.py index 7dca029987e..a076f3daa86 100644 --- a/tests/helpers/test_area_registry.py +++ b/tests/helpers/test_area_registry.py @@ -196,8 +196,9 @@ async def test_load_area(hass, registry): async def test_loading_area_from_storage(hass, hass_storage): """Test loading stored areas on start.""" hass_storage[area_registry.STORAGE_KEY] = { - "version": area_registry.STORAGE_VERSION, - "data": {"areas": [{"id": "12345A", "name": "mock"}]}, + "version": area_registry.STORAGE_VERSION_MAJOR, + "minor_version": area_registry.STORAGE_VERSION_MINOR, + "data": {"areas": [{"id": "12345A", "name": "mock", "picture": "blah"}]}, } await area_registry.async_load(hass) @@ -206,6 +207,31 @@ async def test_loading_area_from_storage(hass, hass_storage): assert len(registry.areas) == 1 +@pytest.mark.parametrize("load_registries", [False]) +async def test_migration_from_1_1(hass, hass_storage): + """Test migration from version 1.1.""" + hass_storage[area_registry.STORAGE_KEY] = { + "version": 1, + "data": {"areas": [{"id": "12345A", "name": "mock"}]}, + } + + await area_registry.async_load(hass) + registry = area_registry.async_get(hass) + + # Test data was loaded + entry = registry.async_get_or_create("mock") + assert entry.id == "12345A" + + # Check we store migrated data + await flush_store(registry._store) + assert hass_storage[area_registry.STORAGE_KEY] == { + "version": area_registry.STORAGE_VERSION_MAJOR, + "minor_version": area_registry.STORAGE_VERSION_MINOR, + "key": area_registry.STORAGE_KEY, + "data": {"areas": [{"id": "12345A", "name": "mock", "picture": None}]}, + } + + async def test_async_get_or_create(hass, registry): """Make sure we can get the area by name.""" area = registry.async_get_or_create("Mock1")