Bump device registry version to 1.2 (#60199)

homeassistant/helpers/device_registry.py

@@ -11,6 +11,7 @@ import attr
 from homeassistant.const import EVENT_HOMEASSISTANT_STARTED
 from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.exceptions import RequiredParameterMissing
+from homeassistant.helpers import storage
 from homeassistant.helpers.frame import report
 from homeassistant.loader import bind_hass
 from homeassistant.util.enum import StrEnum
@@ -31,7 +32,8 @@ _LOGGER = logging.getLogger(__name__)
 DATA_REGISTRY = "device_registry"
 EVENT_DEVICE_REGISTRY_UPDATED = "device_registry_updated"
 STORAGE_KEY = "core.device_registry"
-STORAGE_VERSION = 1
+STORAGE_VERSION_MAJOR = 1
+STORAGE_VERSION_MINOR = 2
 SAVE_DELAY = 10
 CLEANUP_DELAY = 10
 
@@ -159,6 +161,41 @@ def _async_get_device_id_from_index(
     return None
 
 
+class DeviceRegistryStore(storage.Store):
+    """Store entity registry data."""
+
+    async def _async_migrate_func(
+        self, old_major_version: int, old_minor_version: int, old_data: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Migrate to the new version."""
+        if old_major_version < 2 and old_minor_version < 2:
+            # From version 1.1
+            for device in old_data["devices"]:
+                # Introduced in 0.110
+                device["entry_type"] = device.get("entry_type")
+                # Introduced in 0.79
+                # renamed in 0.95
+                device["via_device_id"] = device.get("via_device_id") or device.get(
+                    "hub_device_id"
+                )
+                # Introduced in 0.87
+                device["area_id"] = device.get("area_id")
+                device["name_by_user"] = device.get("name_by_user")
+                # Introduced in 0.119
+                device["disabled_by"] = device.get("disabled_by")
+                # Introduced in 2021.11
+                device["configuration_url"] = device.get("configuration_url")
+            # Introduced in 0.111
+            old_data["deleted_devices"] = old_data.get("deleted_devices", [])
+            for device in old_data["deleted_devices"]:
+                # Introduced in 2021.2
+                device["orphaned_timestamp"] = device.get("orphaned_timestamp")
+
+        if old_major_version > 1:
+            raise NotImplementedError
+        return old_data
+
+
 class DeviceRegistry:
     """Class to hold a registry of devices."""
 
@@ -170,8 +207,12 @@ class DeviceRegistry:
     def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the device registry."""
         self.hass = hass
-        self._store = hass.helpers.storage.Store(
+        self._store = DeviceRegistryStore(
-            STORAGE_VERSION, STORAGE_KEY, atomic_writes=True
+            hass,
+            STORAGE_VERSION_MAJOR,
+            STORAGE_KEY,
+            atomic_writes=True,
+            minor_version=STORAGE_VERSION_MINOR,
         )
         self._clear_index()
 
@@ -519,44 +560,36 @@ class DeviceRegistry:
         deleted_devices = OrderedDict()
 
         if data is not None:
+            data = cast("dict[str, Any]", data)
             for device in data["devices"]:
                 devices[device["id"]] = DeviceEntry(
+                    area_id=device["area_id"],
                     config_entries=set(device["config_entries"]),
+                    configuration_url=device["configuration_url"],
                     # type ignores (if tuple arg was cast): likely https://github.com/python/mypy/issues/8625
                     connections={tuple(conn) for conn in device["connections"]},  # type: ignore[misc]
+                    disabled_by=device["disabled_by"],
+                    entry_type=DeviceEntryType(device["entry_type"])
+                    if device["entry_type"]
+                    else None,
+                    id=device["id"],
                     identifiers={tuple(iden) for iden in device["identifiers"]},  # type: ignore[misc]
                     manufacturer=device["manufacturer"],
                     model=device["model"],
+                    name_by_user=device["name_by_user"],
                     name=device["name"],
                     sw_version=device["sw_version"],
-                    # Introduced in 0.110
+                    via_device_id=device["via_device_id"],
-                    entry_type=DeviceEntryType(device["entry_type"])
-                    if device.get("entry_type")
-                    else None,
-                    id=device["id"],
-                    # Introduced in 0.79
-                    # renamed in 0.95
-                    via_device_id=(
-                        device.get("via_device_id") or device.get("hub_device_id")
-                    ),
-                    # Introduced in 0.87
-                    area_id=device.get("area_id"),
-                    name_by_user=device.get("name_by_user"),
-                    # Introduced in 0.119
-                    disabled_by=device.get("disabled_by"),
-                    # Introduced in 2021.11
-                    configuration_url=device.get("configuration_url"),
                 )
             # Introduced in 0.111
-            for device in data.get("deleted_devices", []):
+            for device in data["deleted_devices"]:
                 deleted_devices[device["id"]] = DeletedDeviceEntry(
                     config_entries=set(device["config_entries"]),
                     # type ignores (if tuple arg was cast): likely https://github.com/python/mypy/issues/8625
                     connections={tuple(conn) for conn in device["connections"]},  # type: ignore[misc]
                     identifiers={tuple(iden) for iden in device["identifiers"]},  # type: ignore[misc]
                     id=device["id"],
-                    # Introduced in 2021.2
+                    orphaned_timestamp=device["orphaned_timestamp"],
-                    orphaned_timestamp=device.get("orphaned_timestamp"),
                 )
 
         self.devices = devices

tests/helpers/test_device_registry.py

@@ -167,23 +167,25 @@ async def test_multiple_config_entries(registry):
 async def test_loading_from_storage(hass, hass_storage):
     """Test loading stored devices on start."""
     hass_storage[device_registry.STORAGE_KEY] = {
-        "version": device_registry.STORAGE_VERSION,
+        "version": device_registry.STORAGE_VERSION_MAJOR,
+        "minor_version": device_registry.STORAGE_VERSION_MINOR,
         "data": {
             "devices": [
                 {
+                    "area_id": "12345A",
                     "config_entries": ["1234"],
+                    "configuration_url": None,
                     "connections": [["Zigbee", "01.23.45.67.89"]],
+                    "disabled_by": device_registry.DISABLED_USER,
+                    "entry_type": device_registry.DeviceEntryType.SERVICE,
                     "id": "abcdefghijklm",
                     "identifiers": [["serial", "12:34:56:AB:CD:EF"]],
                     "manufacturer": "manufacturer",
                     "model": "model",
+                    "name_by_user": "Test Friendly Name",
                     "name": "name",
                     "sw_version": "version",
-                    "entry_type": device_registry.DeviceEntryType.SERVICE,
+                    "via_device_id": None,
-                    "area_id": "12345A",
-                    "name_by_user": "Test Friendly Name",
-                    "disabled_by": device_registry.DISABLED_USER,
-                    "suggested_area": "Kitchen",
                 }
             ],
             "deleted_devices": [
@@ -192,6 +194,7 @@ async def test_loading_from_storage(hass, hass_storage):
                     "connections": [["Zigbee", "23.45.67.89.01"]],
                     "id": "bcdefghijklmn",
                     "identifiers": [["serial", "34:56:AB:CD:EF:12"]],
+                    "orphaned_timestamp": None,
                 }
             ],
         },
@@ -231,6 +234,79 @@ async def test_loading_from_storage(hass, hass_storage):
     assert isinstance(entry.identifiers, set)
 
 
+@pytest.mark.parametrize("load_registries", [False])
+async def test_migration_1_1_to_1_2(hass, hass_storage):
+    """Test migration from version 1.1 to 1.2."""
+    hass_storage[device_registry.STORAGE_KEY] = {
+        "version": 1,
+        "minor_version": 1,
+        "data": {
+            "devices": [
+                {
+                    "config_entries": ["1234"],
+                    "connections": [["Zigbee", "01.23.45.67.89"]],
+                    "entry_type": "service",
+                    "id": "abcdefghijklm",
+                    "identifiers": [["serial", "12:34:56:AB:CD:EF"]],
+                    "manufacturer": "manufacturer",
+                    "model": "model",
+                    "name": "name",
+                    "sw_version": "version",
+                }
+            ],
+        },
+    }
+
+    await device_registry.async_load(hass)
+    registry = device_registry.async_get(hass)
+
+    # Test data was loaded
+    entry = registry.async_get_or_create(
+        config_entry_id="1234",
+        connections={("Zigbee", "01.23.45.67.89")},
+        identifiers={("serial", "12:34:56:AB:CD:EF")},
+    )
+    assert entry.id == "abcdefghijklm"
+
+    # Update to trigger a store
+    entry = registry.async_get_or_create(
+        config_entry_id="1234",
+        connections={("Zigbee", "01.23.45.67.89")},
+        identifiers={("serial", "12:34:56:AB:CD:EF")},
+        sw_version="new_version",
+    )
+    assert entry.id == "abcdefghijklm"
+
+    # Check we store migrated data
+    await flush_store(registry._store)
+    assert hass_storage[device_registry.STORAGE_KEY] == {
+        "version": device_registry.STORAGE_VERSION_MAJOR,
+        "minor_version": device_registry.STORAGE_VERSION_MINOR,
+        "key": device_registry.STORAGE_KEY,
+        "data": {
+            "devices": [
+                {
+                    "area_id": None,
+                    "config_entries": ["1234"],
+                    "configuration_url": None,
+                    "connections": [["Zigbee", "01.23.45.67.89"]],
+                    "disabled_by": None,
+                    "entry_type": "service",
+                    "id": "abcdefghijklm",
+                    "identifiers": [["serial", "12:34:56:AB:CD:EF"]],
+                    "manufacturer": "manufacturer",
+                    "model": "model",
+                    "name": "name",
+                    "name_by_user": None,
+                    "sw_version": "new_version",
+                    "via_device_id": None,
+                }
+            ],
+            "deleted_devices": [],
+        },
+    }
+
+
 async def test_removing_config_entries(hass, registry, update_events):
     """Make sure we do not get duplicate entries."""
     entry = registry.async_get_or_create(