Add minor version to area registry store (#84280)

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]]]]](
+        self._store = AreaRegistryStore(
-            hass, STORAGE_VERSION, STORAGE_KEY, atomic_writes=True
+            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
+                    name=area["name"],
-                    picture=area.get("picture"),
                     normalized_name=normalized_name,
+                    picture=area["picture"],
                 )
                 self._normalized_name_area_idx[normalized_name] = area["id"]
 

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,
+        "version": area_registry.STORAGE_VERSION_MAJOR,
-        "data": {"areas": [{"id": "12345A", "name": "mock"}]},
+        "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")