Improve migration to device registry version 1.11 (#151315)

Co-authored-by: Franck Nijhof <git@frenck.dev>

homeassistant/helpers/device_registry.py

@@ -9,7 +9,7 @@ from enum import StrEnum
 from functools import lru_cache
 import logging
 import time
-from typing import TYPE_CHECKING, Any, Literal, TypedDict
+from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
 
 import attr
 from yarl import URL
@@ -68,6 +68,8 @@ CONNECTION_ZIGBEE = "zigbee"
 
 ORPHANED_DEVICE_KEEP_SECONDS = 86400 * 30
 
+UNDEFINED_STR: Final = "UNDEFINED"
+
 # Can be removed when suggested_area is removed from DeviceEntry
 RUNTIME_ONLY_ATTRS = {"suggested_area"}
 
@@ -463,7 +465,7 @@ class DeletedDeviceEntry:
         validator=_normalize_connections_validator
     )
     created_at: datetime = attr.ib()
-    disabled_by: DeviceEntryDisabler | None = attr.ib()
+    disabled_by: DeviceEntryDisabler | UndefinedType | None = attr.ib()
     id: str = attr.ib()
     identifiers: set[tuple[str, str]] = attr.ib()
     labels: set[str] = attr.ib()
@@ -478,15 +480,19 @@ class DeletedDeviceEntry:
         config_subentry_id: str | None,
         connections: set[tuple[str, str]],
         identifiers: set[tuple[str, str]],
+        disabled_by: DeviceEntryDisabler | UndefinedType | None,
     ) -> DeviceEntry:
         """Create DeviceEntry from DeletedDeviceEntry."""
         # Adjust disabled_by based on config entry state
-        disabled_by = self.disabled_by
+        if self.disabled_by is not UNDEFINED:
-        if config_entry.disabled_by:
+            disabled_by = self.disabled_by
-            if disabled_by is None:
+            if config_entry.disabled_by:
-                disabled_by = DeviceEntryDisabler.CONFIG_ENTRY
+                if disabled_by is None:
-        elif disabled_by == DeviceEntryDisabler.CONFIG_ENTRY:
+                    disabled_by = DeviceEntryDisabler.CONFIG_ENTRY
-            disabled_by = None
+            elif disabled_by == DeviceEntryDisabler.CONFIG_ENTRY:
+                disabled_by = None
+        else:
+            disabled_by = disabled_by if disabled_by is not UNDEFINED else None
         return DeviceEntry(
             area_id=self.area_id,
             # type ignores: likely https://github.com/python/mypy/issues/8625
@@ -517,7 +523,9 @@ class DeletedDeviceEntry:
                     },
                     "connections": list(self.connections),
                     "created_at": self.created_at,
-                    "disabled_by": self.disabled_by,
+                    "disabled_by": self.disabled_by
+                    if self.disabled_by is not UNDEFINED
+                    else UNDEFINED_STR,
                     "identifiers": list(self.identifiers),
                     "id": self.id,
                     "labels": list(self.labels),
@@ -605,7 +613,7 @@ class DeviceRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]):
                 # Introduced in 2025.6
                 for device in old_data["deleted_devices"]:
                     device["area_id"] = None
-                    device["disabled_by"] = None
+                    device["disabled_by"] = UNDEFINED_STR
                     device["labels"] = []
                     device["name_by_user"] = None
             if old_minor_version < 11:
@@ -935,6 +943,7 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
                     config_subentry_id if config_subentry_id is not UNDEFINED else None,
                     connections,
                     identifiers,
+                    disabled_by,
                 )
                 disabled_by = UNDEFINED
 
@@ -1502,7 +1511,21 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
                     sw_version=device["sw_version"],
                     via_device_id=device["via_device_id"],
                 )
+
             # Introduced in 0.111
+            def get_optional_enum[_EnumT: StrEnum](
+                cls: type[_EnumT], value: str | None
+            ) -> _EnumT | UndefinedType | None:
+                """Convert string to the passed enum, UNDEFINED or None."""
+                if value is None:
+                    return None
+                if value == UNDEFINED_STR:
+                    return UNDEFINED
+                try:
+                    return cls(value)
+                except ValueError:
+                    return None
+
             for device in data["deleted_devices"]:
                 deleted_devices[device["id"]] = DeletedDeviceEntry(
                     area_id=device["area_id"],
@@ -1515,10 +1538,8 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
                     },
                     connections={tuple(conn) for conn in device["connections"]},
                     created_at=datetime.fromisoformat(device["created_at"]),
-                    disabled_by=(
+                    disabled_by=get_optional_enum(
-                        DeviceEntryDisabler(device["disabled_by"])
+                        DeviceEntryDisabler, device["disabled_by"]
-                        if device["disabled_by"]
-                        else None
                     ),
                     identifiers={tuple(iden) for iden in device["identifiers"]},
                     id=device["id"],

tests/helpers/test_device_registry.py

@@ -8,6 +8,7 @@ import time
 from typing import Any
 from unittest.mock import ANY, patch
 
+import attr
 from freezegun.api import FrozenDateTimeFactory
 import pytest
 from yarl import URL
@@ -21,6 +22,7 @@ from homeassistant.helpers import (
     device_registry as dr,
     entity_registry as er,
 )
+from homeassistant.helpers.typing import UNDEFINED, UndefinedType
 from homeassistant.util.dt import utcnow
 
 from tests.common import MockConfigEntry, async_capture_events, flush_store
@@ -508,6 +510,9 @@ async def test_migration_from_1_1(
     )
     assert entry.id == "abcdefghijklm"
 
+    deleted_entry = registry.deleted_devices["deletedid"]
+    assert deleted_entry.disabled_by is UNDEFINED
+
     # Update to trigger a store
     entry = registry.async_get_or_create(
         config_entry_id=mock_config_entry.entry_id,
@@ -581,7 +586,7 @@ async def test_migration_from_1_1(
                     "config_entries_subentries": {"123456": [None]},
                     "connections": [],
                     "created_at": "1970-01-01T00:00:00+00:00",
-                    "disabled_by": None,
+                    "disabled_by": "UNDEFINED",
                     "id": "deletedid",
                     "identifiers": [["serial", "123456ABCDFF"]],
                     "labels": [],
@@ -3833,6 +3838,130 @@ async def test_restore_device(
     }
 
 
+@pytest.mark.parametrize(
+    ("device_disabled_by", "expected_disabled_by"),
+    [
+        (None, None),
+        (dr.DeviceEntryDisabler.CONFIG_ENTRY, dr.DeviceEntryDisabler.CONFIG_ENTRY),
+        (dr.DeviceEntryDisabler.INTEGRATION, dr.DeviceEntryDisabler.INTEGRATION),
+        (dr.DeviceEntryDisabler.USER, dr.DeviceEntryDisabler.USER),
+        (UNDEFINED, None),
+    ],
+)
+@pytest.mark.usefixtures("freezer")
+async def test_restore_migrated_device_disabled_by(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mock_config_entry: MockConfigEntry,
+    device_disabled_by: dr.DeviceEntryDisabler | UndefinedType | None,
+    expected_disabled_by: dr.DeviceEntryDisabler | None,
+) -> None:
+    """Check how the disabled_by flag is treated when restoring a device."""
+    entry_id = mock_config_entry.entry_id
+    update_events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED)
+    entry = device_registry.async_get_or_create(
+        config_entry_id=entry_id,
+        config_subentry_id=None,
+        configuration_url="http://config_url_orig.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        disabled_by=None,
+        entry_type=dr.DeviceEntryType.SERVICE,
+        hw_version="hw_version_orig",
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer_orig",
+        model="model_orig",
+        model_id="model_id_orig",
+        name="name_orig",
+        serial_number="serial_no_orig",
+        suggested_area="suggested_area_orig",
+        sw_version="version_orig",
+        via_device="via_device_id_orig",
+    )
+
+    assert len(device_registry.devices) == 1
+    assert len(device_registry.deleted_devices) == 0
+
+    device_registry.async_remove_device(entry.id)
+
+    assert len(device_registry.devices) == 0
+    assert len(device_registry.deleted_devices) == 1
+
+    deleted_entry = device_registry.deleted_devices[entry.id]
+    device_registry.deleted_devices[entry.id] = attr.evolve(
+        deleted_entry, disabled_by=UNDEFINED
+    )
+
+    # This will restore the original device, user customizations of
+    # area_id, disabled_by, labels and name_by_user will be restored
+    entry3 = device_registry.async_get_or_create(
+        config_entry_id=entry_id,
+        config_subentry_id=None,
+        configuration_url="http://config_url_new.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        disabled_by=device_disabled_by,
+        entry_type=None,
+        hw_version="hw_version_new",
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer_new",
+        model="model_new",
+        model_id="model_id_new",
+        name="name_new",
+        serial_number="serial_no_new",
+        suggested_area="suggested_area_new",
+        sw_version="version_new",
+        via_device="via_device_id_new",
+    )
+    assert entry3 == dr.DeviceEntry(
+        area_id="suggested_area_orig",
+        config_entries={entry_id},
+        config_entries_subentries={entry_id: {None}},
+        configuration_url="http://config_url_new.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")},
+        created_at=utcnow(),
+        disabled_by=expected_disabled_by,
+        entry_type=None,
+        hw_version="hw_version_new",
+        id=entry.id,
+        identifiers={("bridgeid", "0123")},
+        labels=set(),
+        manufacturer="manufacturer_new",
+        model="model_new",
+        model_id="model_id_new",
+        modified_at=utcnow(),
+        name_by_user=None,
+        name="name_new",
+        primary_config_entry=entry_id,
+        serial_number="serial_no_new",
+        suggested_area="suggested_area_new",
+        sw_version="version_new",
+    )
+
+    assert entry.id == entry3.id
+    assert len(device_registry.devices) == 1
+    assert len(device_registry.deleted_devices) == 0
+
+    assert isinstance(entry3.config_entries, set)
+    assert isinstance(entry3.connections, set)
+    assert isinstance(entry3.identifiers, set)
+
+    await hass.async_block_till_done()
+
+    assert len(update_events) == 3
+    assert update_events[0].data == {
+        "action": "create",
+        "device_id": entry.id,
+    }
+    assert update_events[1].data == {
+        "action": "remove",
+        "device_id": entry.id,
+        "device": entry.dict_repr,
+    }
+    assert update_events[2].data == {
+        "action": "create",
+        "device_id": entry3.id,
+    }
+
+
 @pytest.mark.parametrize(
     (
         "config_entry_disabled_by",