Improve migration to device registry version 1.10 (#151571)

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
@@ -349,6 +351,7 @@ async def test_loading_from_storage(
                     "connections": [["Zigbee", "23.45.67.89.01"]],
                     "created_at": created_at,
                     "disabled_by": dr.DeviceEntryDisabler.USER,
+                    "disabled_by_undefined": False,
                     "id": "bcdefghijklmn",
                     "identifiers": [["serial", "3456ABCDEF12"]],
                     "labels": {"label1", "label2"},
@@ -508,6 +511,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,
@@ -582,6 +588,7 @@ async def test_migration_from_1_1(
                     "connections": [],
                     "created_at": "1970-01-01T00:00:00+00:00",
                     "disabled_by": None,
+                    "disabled_by_undefined": True,
                     "id": "deletedid",
                     "identifiers": [["serial", "123456ABCDFF"]],
                     "labels": [],
@@ -1477,6 +1484,7 @@ async def test_migration_from_1_10(
                     "connections": [["mac", "123456ABCDAB"]],
                     "created_at": "1970-01-01T00:00:00+00:00",
                     "disabled_by": None,
+                    "disabled_by_undefined": False,
                     "id": "abcdefghijklm2",
                     "identifiers": [["serial", "123456ABCDAB"]],
                     "labels": [],
@@ -1553,6 +1561,144 @@ async def test_migration_from_1_10(
                     "connections": [["mac", "12:34:56:ab:cd:ab"]],
                     "created_at": "1970-01-01T00:00:00+00:00",
                     "disabled_by": None,
+                    "disabled_by_undefined": False,
+                    "id": "abcdefghijklm2",
+                    "identifiers": [["serial", "123456ABCDAB"]],
+                    "labels": [],
+                    "modified_at": "1970-01-01T00:00:00+00:00",
+                    "name_by_user": None,
+                    "orphaned_timestamp": "1970-01-01T00:00:00+00:00",
+                },
+            ],
+        },
+    }
+
+
+@pytest.mark.parametrize("load_registries", [False])
+@pytest.mark.usefixtures("freezer")
+async def test_migration_from_1_11(
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test migration from version 1.11."""
+    hass_storage[dr.STORAGE_KEY] = {
+        "version": 1,
+        "minor_version": 10,
+        "key": dr.STORAGE_KEY,
+        "data": {
+            "devices": [
+                {
+                    "area_id": None,
+                    "config_entries": [mock_config_entry.entry_id],
+                    "config_entries_subentries": {mock_config_entry.entry_id: [None]},
+                    "configuration_url": None,
+                    "connections": [["mac", "123456ABCDEF"]],
+                    "created_at": "1970-01-01T00:00:00+00:00",
+                    "disabled_by": None,
+                    "entry_type": "service",
+                    "hw_version": "hw_version",
+                    "id": "abcdefghijklm",
+                    "identifiers": [["serial", "123456ABCDEF"]],
+                    "labels": ["blah"],
+                    "manufacturer": "manufacturer",
+                    "model": "model",
+                    "name": "name",
+                    "model_id": None,
+                    "modified_at": "1970-01-01T00:00:00+00:00",
+                    "name_by_user": None,
+                    "primary_config_entry": mock_config_entry.entry_id,
+                    "serial_number": None,
+                    "sw_version": "new_version",
+                    "via_device_id": None,
+                },
+            ],
+            "deleted_devices": [
+                {
+                    "area_id": None,
+                    "config_entries": ["234567"],
+                    "config_entries_subentries": {"234567": [None]},
+                    "connections": [["mac", "123456ABCDAB"]],
+                    "created_at": "1970-01-01T00:00:00+00:00",
+                    "disabled_by": None,
+                    "disabled_by_undefined": False,
+                    "id": "abcdefghijklm2",
+                    "identifiers": [["serial", "123456ABCDAB"]],
+                    "labels": [],
+                    "modified_at": "1970-01-01T00:00:00+00:00",
+                    "name_by_user": None,
+                    "orphaned_timestamp": "1970-01-01T00:00:00+00:00",
+                },
+            ],
+        },
+    }
+
+    await dr.async_load(hass)
+    registry = dr.async_get(hass)
+
+    # Test data was loaded
+    entry = registry.async_get_or_create(
+        config_entry_id=mock_config_entry.entry_id,
+        identifiers={("serial", "123456ABCDEF")},
+    )
+    assert entry.id == "abcdefghijklm"
+    deleted_entry = registry.deleted_devices.get_entry(
+        connections=set(),
+        identifiers={("serial", "123456ABCDAB")},
+    )
+    assert deleted_entry.id == "abcdefghijklm2"
+
+    # Update to trigger a store
+    entry = registry.async_get_or_create(
+        config_entry_id=mock_config_entry.entry_id,
+        identifiers={("serial", "123456ABCDEF")},
+        sw_version="new_version",
+    )
+    assert entry.id == "abcdefghijklm"
+
+    # Check we store migrated data
+    await flush_store(registry._store)
+
+    assert hass_storage[dr.STORAGE_KEY] == {
+        "version": dr.STORAGE_VERSION_MAJOR,
+        "minor_version": dr.STORAGE_VERSION_MINOR,
+        "key": dr.STORAGE_KEY,
+        "data": {
+            "devices": [
+                {
+                    "area_id": None,
+                    "config_entries": [mock_config_entry.entry_id],
+                    "config_entries_subentries": {mock_config_entry.entry_id: [None]},
+                    "configuration_url": None,
+                    "connections": [["mac", "12:34:56:ab:cd:ef"]],
+                    "created_at": "1970-01-01T00:00:00+00:00",
+                    "disabled_by": None,
+                    "entry_type": "service",
+                    "hw_version": "hw_version",
+                    "id": "abcdefghijklm",
+                    "identifiers": [["serial", "123456ABCDEF"]],
+                    "labels": ["blah"],
+                    "manufacturer": "manufacturer",
+                    "model": "model",
+                    "name": "name",
+                    "model_id": None,
+                    "modified_at": "1970-01-01T00:00:00+00:00",
+                    "name_by_user": None,
+                    "primary_config_entry": mock_config_entry.entry_id,
+                    "serial_number": None,
+                    "sw_version": "new_version",
+                    "via_device_id": None,
+                },
+            ],
+            "deleted_devices": [
+                {
+                    "area_id": None,
+                    "config_entries": ["234567"],
+                    "config_entries_subentries": {"234567": [None]},
+                    "connections": [["mac", "12:34:56:ab:cd:ab"]],
+                    "created_at": "1970-01-01T00:00:00+00:00",
+                    "disabled_by": None,
+                    "disabled_by_undefined": False,
                     "id": "abcdefghijklm2",
                     "identifiers": [["serial", "123456ABCDAB"]],
                     "labels": [],
@@ -3833,6 +3979,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",