Add created_at/modified_at to floor registry (#122071)

homeassistant/helpers/floor_registry.py

@@ -5,10 +5,12 @@ from __future__ import annotations
 from collections.abc import Iterable
 import dataclasses
 from dataclasses import dataclass
-from typing import Literal, TypedDict
+from datetime import datetime
+from typing import Any, Literal, TypedDict
 
 from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.util import slugify
+from homeassistant.util.dt import utc_from_timestamp, utcnow
 from homeassistant.util.event_type import EventType
 from homeassistant.util.hass_dict import HassKey
 
@@ -28,6 +30,7 @@ EVENT_FLOOR_REGISTRY_UPDATED: EventType[EventFloorRegistryUpdatedData] = EventTy
 )
 STORAGE_KEY = "core.floor_registry"
 STORAGE_VERSION_MAJOR = 1
+STORAGE_VERSION_MINOR = 2
 
 
 class _FloorStoreData(TypedDict):
@@ -38,6 +41,8 @@ class _FloorStoreData(TypedDict):
     icon: str | None
     level: int | None
     name: str
+    created_at: datetime
+    modified_at: datetime
 
 
 class FloorRegistryStoreData(TypedDict):
@@ -66,6 +71,28 @@ class FloorEntry(NormalizedNameBaseRegistryEntry):
     level: int | None = None
 
 
+class FloorRegistryStore(Store[FloorRegistryStoreData]):
+    """Store floor registry data."""
+
+    async def _async_migrate_func(
+        self,
+        old_major_version: int,
+        old_minor_version: int,
+        old_data: dict[str, list[dict[str, Any]]],
+    ) -> FloorRegistryStoreData:
+        """Migrate to the new version."""
+        if old_major_version > STORAGE_VERSION_MAJOR:
+            raise ValueError("Can't migrate to future version")
+
+        if old_major_version == 1:
+            if old_minor_version < 2:
+                # Version 1.2 implements migration and adds created_at and modified_at
+                for floor in old_data["floors"]:
+                    floor["created_at"] = floor["modified_at"] = utc_from_timestamp(0)
+
+        return old_data  # type: ignore[return-value]
+
+
 class FloorRegistry(BaseRegistry[FloorRegistryStoreData]):
     """Class to hold a registry of floors."""
 
@@ -75,11 +102,12 @@ class FloorRegistry(BaseRegistry[FloorRegistryStoreData]):
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the floor registry."""
         self.hass = hass
-        self._store = Store(
+        self._store = FloorRegistryStore(
             hass,
             STORAGE_VERSION_MAJOR,
             STORAGE_KEY,
             atomic_writes=True,
+            minor_version=STORAGE_VERSION_MINOR,
         )
 
     @callback
@@ -175,7 +203,7 @@ class FloorRegistry(BaseRegistry[FloorRegistryStoreData]):
     ) -> FloorEntry:
         """Update name of the floor."""
         old = self.floors[floor_id]
-        changes = {
+        changes: dict[str, Any] = {
             attr_name: value
             for attr_name, value in (
                 ("aliases", aliases),
@@ -191,8 +219,10 @@ class FloorRegistry(BaseRegistry[FloorRegistryStoreData]):
         if not changes:
             return old
 
+        changes["modified_at"] = utcnow()
+
         self.hass.verify_event_loop_thread("floor_registry.async_update")
-        new = self.floors[floor_id] = dataclasses.replace(old, **changes)  # type: ignore[arg-type]
+        new = self.floors[floor_id] = dataclasses.replace(old, **changes)
 
         self.async_schedule_save()
         self.hass.bus.async_fire_internal(
@@ -220,6 +250,8 @@ class FloorRegistry(BaseRegistry[FloorRegistryStoreData]):
                     name=floor["name"],
                     level=floor["level"],
                     normalized_name=normalized_name,
+                    created_at=floor["created_at"],
+                    modified_at=floor["modified_at"],
                 )
 
         self.floors = floors
@@ -236,6 +268,8 @@ class FloorRegistry(BaseRegistry[FloorRegistryStoreData]):
                     "icon": entry.icon,
                     "level": entry.level,
                     "name": entry.name,
+                    "created_at": entry.created_at,
+                    "modified_at": entry.modified_at,
                 }
                 for entry in self.floors.values()
             ]

tests/helpers/test_floor_registry.py

@@ -1,15 +1,18 @@
 """Tests for the floor registry."""
 
+from datetime import datetime
 from functools import partial
 import re
 from typing import Any
 
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import area_registry as ar, floor_registry as fr
+from homeassistant.util.dt import utcnow
 
-from tests.common import async_capture_events, flush_store
+from tests.common import ANY, async_capture_events, flush_store
 
 
 async def test_list_floors(floor_registry: fr.FloorRegistry) -> None:
@@ -18,8 +21,10 @@ async def test_list_floors(floor_registry: fr.FloorRegistry) -> None:
     assert len(list(floors)) == len(floor_registry.floors)
 
 
+@pytest.mark.usefixtures("freezer")
 async def test_create_floor(
-    hass: HomeAssistant, floor_registry: fr.FloorRegistry
+    hass: HomeAssistant,
+    floor_registry: fr.FloorRegistry,
 ) -> None:
     """Make sure that we can create floors."""
     update_events = async_capture_events(hass, fr.EVENT_FLOOR_REGISTRY_UPDATED)
@@ -30,11 +35,16 @@ async def test_create_floor(
         level=1,
     )
 
-    assert floor.floor_id == "first_floor"
-    assert floor.name == "First floor"
-    assert floor.icon == "mdi:home-floor-1"
-    assert floor.aliases == {"first", "ground", "ground floor"}
-    assert floor.level == 1
+    assert floor == fr.FloorEntry(
+        floor_id="first_floor",
+        name="First floor",
+        icon="mdi:home-floor-1",
+        aliases={"first", "ground", "ground floor"},
+        level=1,
+        created_at=utcnow(),
+        modified_at=utcnow(),
+        normalized_name=ANY,
+    )
 
     assert len(floor_registry.floors) == 1
 
@@ -116,18 +126,31 @@ async def test_delete_non_existing_floor(floor_registry: fr.FloorRegistry) -> No
 
 
 async def test_update_floor(
-    hass: HomeAssistant, floor_registry: fr.FloorRegistry
+    hass: HomeAssistant,
+    floor_registry: fr.FloorRegistry,
+    freezer: FrozenDateTimeFactory,
 ) -> None:
     """Make sure that we can update floors."""
+    created_at = datetime.fromisoformat("2024-01-01T01:00:00+00:00")
+    freezer.move_to(created_at)
+
     update_events = async_capture_events(hass, fr.EVENT_FLOOR_REGISTRY_UPDATED)
     floor = floor_registry.async_create("First floor")
 
+    assert floor == fr.FloorEntry(
+        floor_id="first_floor",
+        name="First floor",
+        icon=None,
+        aliases=set(),
+        level=None,
+        created_at=created_at,
+        modified_at=created_at,
+        normalized_name=ANY,
+    )
     assert len(floor_registry.floors) == 1
-    assert floor.floor_id == "first_floor"
-    assert floor.name == "First floor"
-    assert floor.icon is None
-    assert floor.aliases == set()
-    assert floor.level is None
+
+    modified_at = datetime.fromisoformat("2024-02-01T01:00:00+00:00")
+    freezer.move_to(modified_at)
 
     updated_floor = floor_registry.async_update(
         floor.floor_id,
@@ -138,11 +161,16 @@ async def test_update_floor(
     )
 
     assert updated_floor != floor
-    assert updated_floor.floor_id == "first_floor"
-    assert updated_floor.name == "Second floor"
-    assert updated_floor.icon == "mdi:home-floor-2"
-    assert updated_floor.aliases == {"ground", "downstairs"}
-    assert updated_floor.level == 2
+    assert updated_floor == fr.FloorEntry(
+        floor_id="first_floor",
+        name="Second floor",
+        icon="mdi:home-floor-2",
+        aliases={"ground", "downstairs"},
+        level=2,
+        created_at=created_at,
+        modified_at=modified_at,
+        normalized_name=ANY,
+    )
 
     assert len(floor_registry.floors) == 1
 
@@ -280,7 +308,8 @@ async def test_load_floors(
 
 @pytest.mark.parametrize("load_registries", [False])
 async def test_loading_floors_from_storage(
-    hass: HomeAssistant, hass_storage: dict[str, Any]
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
 ) -> None:
     """Test loading stored floors on start."""
     hass_storage[fr.STORAGE_KEY] = {
@@ -392,3 +421,52 @@ async def test_async_update_thread_safety(
         await hass.async_add_executor_job(
             partial(floor_registry.async_update, any_floor.floor_id, name="new name")
         )
+
+
+@pytest.mark.parametrize("load_registries", [False])
+async def test_migration_from_1_1(
+    hass: HomeAssistant, hass_storage: dict[str, Any]
+) -> None:
+    """Test migration from version 1.1."""
+    hass_storage[fr.STORAGE_KEY] = {
+        "version": 1,
+        "data": {
+            "floors": [
+                {
+                    "floor_id": "12345A",
+                    "name": "mock",
+                    "aliases": [],
+                    "icon": None,
+                    "level": None,
+                }
+            ]
+        },
+    }
+
+    await fr.async_load(hass)
+    registry = fr.async_get(hass)
+
+    # Test data was loaded
+    entry = registry.async_get_floor_by_name("mock")
+    assert entry.floor_id == "12345A"
+
+    # Check we store migrated data
+    await flush_store(registry._store)
+    assert hass_storage[fr.STORAGE_KEY] == {
+        "version": fr.STORAGE_VERSION_MAJOR,
+        "minor_version": fr.STORAGE_VERSION_MINOR,
+        "key": fr.STORAGE_KEY,
+        "data": {
+            "floors": [
+                {
+                    "aliases": [],
+                    "icon": None,
+                    "floor_id": "12345A",
+                    "level": None,
+                    "name": "mock",
+                    "created_at": "1970-01-01T00:00:00+00:00",
+                    "modified_at": "1970-01-01T00:00:00+00:00",
+                }
+            ]
+        },
+    }