Improve device registry restore tests (#145186)

tests/helpers/test_device_registry.py

@@ -6,7 +6,7 @@ from datetime import datetime
 from functools import partial
 import time
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import ANY, patch
 
 from freezegun.api import FrozenDateTimeFactory
 import pytest
@@ -34,6 +34,32 @@ def mock_config_entry(hass: HomeAssistant) -> MockConfigEntry:
     return entry
 
 
+@pytest.fixture
+def mock_config_entry_with_subentries(hass: HomeAssistant) -> MockConfigEntry:
+    """Create a mock config entry and add it to hass."""
+    entry = MockConfigEntry(
+        title=None,
+        subentries_data=(
+            config_entries.ConfigSubentryData(
+                data={},
+                subentry_id="mock-subentry-id-1-1",
+                subentry_type="test",
+                title="Mock title",
+                unique_id="test",
+            ),
+            config_entries.ConfigSubentryData(
+                data={},
+                subentry_id="mock-subentry-id-1-2",
+                subentry_type="test",
+                title="Mock title",
+                unique_id="test",
+            ),
+        ),
+    )
+    entry.add_to_hass(hass)
+    return entry
+
+
 async def test_get_or_create_returns_same_entry(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
@@ -3173,19 +3199,41 @@ async def test_cleanup_entity_registry_change(
         assert len(mock_call.mock_calls) == 2
 
 
+@pytest.mark.usefixtures("freezer")
 async def test_restore_device(
     hass: HomeAssistant,
     device_registry: dr.DeviceRegistry,
-    mock_config_entry: MockConfigEntry,
+    mock_config_entry_with_subentries: MockConfigEntry,
 ) -> None:
     """Make sure device id is stable."""
+    entry_id = mock_config_entry_with_subentries.entry_id
+    subentry_id = "mock-subentry-id-1-1"
     update_events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED)
     entry = device_registry.async_get_or_create(
-        config_entry_id=mock_config_entry.entry_id,
+        config_entry_id=entry_id,
+        config_subentry_id=subentry_id,
+        configuration_url="http://config_url_orig.bla",
         connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        entry_type=dr.DeviceEntryType.SERVICE,
+        hw_version="hw_version_orig",
         identifiers={("bridgeid", "0123")},
-        manufacturer="manufacturer",
+        manufacturer="manufacturer_orig",
-        model="model",
+        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",
+    )
+
+    # Apply user customizations
+    device_registry.async_update_device(
+        entry.id,
+        area_id="12345A",
+        disabled_by=dr.DeviceEntryDisabler.USER,
+        labels={"label1", "label2"},
+        name_by_user="Test Friendly Name",
     )
 
     assert len(device_registry.devices) == 1
@@ -3196,19 +3244,79 @@ async def test_restore_device(
     assert len(device_registry.devices) == 0
     assert len(device_registry.deleted_devices) == 1
 
+    # This will create a new device
     entry2 = device_registry.async_get_or_create(
-        config_entry_id=mock_config_entry.entry_id,
+        config_entry_id=entry_id,
         connections={(dr.CONNECTION_NETWORK_MAC, "34:56:78:CD:EF:12")},
         identifiers={("bridgeid", "4567")},
         manufacturer="manufacturer",
         model="model",
     )
-    entry3 = device_registry.async_get_or_create(
+    assert entry2 == dr.DeviceEntry(
-        config_entry_id=mock_config_entry.entry_id,
+        area_id=None,
-        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        config_entries={entry_id},
-        identifiers={("bridgeid", "0123")},
+        config_entries_subentries={entry_id: {None}},
+        configuration_url=None,
+        connections={(dr.CONNECTION_NETWORK_MAC, "34:56:78:cd:ef:12")},
+        created_at=utcnow(),
+        disabled_by=None,
+        entry_type=None,
+        hw_version=None,
+        id=ANY,
+        identifiers={("bridgeid", "4567")},
+        labels={},
         manufacturer="manufacturer",
         model="model",
+        model_id=None,
+        modified_at=utcnow(),
+        name_by_user=None,
+        name=None,
+        primary_config_entry=entry_id,
+        serial_number=None,
+        suggested_area=None,
+        sw_version=None,
+    )
+    # This will restore the original device
+    entry3 = device_registry.async_get_or_create(
+        config_entry_id=entry_id,
+        config_subentry_id=subentry_id,
+        configuration_url="http://config_url_new.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        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_new",
+        config_entries={entry_id},
+        config_entries_subentries={entry_id: {subentry_id}},
+        configuration_url="http://config_url_new.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")},
+        created_at=utcnow(),
+        disabled_by=None,
+        entry_type=None,
+        hw_version="hw_version_new",
+        id=entry.id,
+        identifiers={("bridgeid", "0123")},
+        labels={},
+        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
@@ -3222,129 +3330,186 @@ async def test_restore_device(
 
     await hass.async_block_till_done()
 
-    assert len(update_events) == 4
+    assert len(update_events) == 5
     assert update_events[0].data == {
         "action": "create",
         "device_id": entry.id,
     }
     assert update_events[1].data == {
-        "action": "remove",
+        "action": "update",
+        "changes": {
+            "area_id": "suggested_area_orig",
+            "disabled_by": None,
+            "labels": set(),
+            "name_by_user": None,
+        },
         "device_id": entry.id,
     }
     assert update_events[2].data == {
+        "action": "remove",
+        "device_id": entry.id,
+    }
+    assert update_events[3].data == {
         "action": "create",
         "device_id": entry2.id,
     }
-    assert update_events[3].data == {
+    assert update_events[4].data == {
-        "action": "create",
-        "device_id": entry3.id,
-    }
-
-
-async def test_restore_simple_device(
-    hass: HomeAssistant,
-    device_registry: dr.DeviceRegistry,
-    mock_config_entry: MockConfigEntry,
-) -> None:
-    """Make sure device id is stable."""
-    update_events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED)
-    entry = device_registry.async_get_or_create(
-        config_entry_id=mock_config_entry.entry_id,
-        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
-        identifiers={("bridgeid", "0123")},
-    )
-
-    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
-
-    entry2 = device_registry.async_get_or_create(
-        config_entry_id=mock_config_entry.entry_id,
-        connections={(dr.CONNECTION_NETWORK_MAC, "34:56:78:CD:EF:12")},
-        identifiers={("bridgeid", "4567")},
-    )
-    entry3 = device_registry.async_get_or_create(
-        config_entry_id=mock_config_entry.entry_id,
-        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
-        identifiers={("bridgeid", "0123")},
-    )
-
-    assert entry.id == entry3.id
-    assert entry.id != entry2.id
-    assert len(device_registry.devices) == 2
-    assert len(device_registry.deleted_devices) == 0
-
-    await hass.async_block_till_done()
-
-    assert len(update_events) == 4
-    assert update_events[0].data == {
-        "action": "create",
-        "device_id": entry.id,
-    }
-    assert update_events[1].data == {
-        "action": "remove",
-        "device_id": entry.id,
-    }
-    assert update_events[2].data == {
-        "action": "create",
-        "device_id": entry2.id,
-    }
-    assert update_events[3].data == {
         "action": "create",
         "device_id": entry3.id,
     }
 
 
+@pytest.mark.usefixtures("freezer")
 async def test_restore_shared_device(
     hass: HomeAssistant, device_registry: dr.DeviceRegistry
 ) -> None:
     """Make sure device id is stable for shared devices."""
     update_events = async_capture_events(hass, dr.EVENT_DEVICE_REGISTRY_UPDATED)
-    config_entry_1 = MockConfigEntry()
+    config_entry_1 = MockConfigEntry(
+        subentries_data=(
+            config_entries.ConfigSubentryData(
+                data={},
+                subentry_id="mock-subentry-id-1-1",
+                subentry_type="test",
+                title="Mock title",
+                unique_id="test",
+            ),
+        ),
+    )
     config_entry_1.add_to_hass(hass)
     config_entry_2 = MockConfigEntry()
     config_entry_2.add_to_hass(hass)
 
     entry = device_registry.async_get_or_create(
         config_entry_id=config_entry_1.entry_id,
+        config_subentry_id="mock-subentry-id-1-1",
+        configuration_url="http://config_url_orig_1.bla",
         connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        entry_type=dr.DeviceEntryType.SERVICE,
+        hw_version="hw_version_orig_1",
         identifiers={("entry_123", "0123")},
-        manufacturer="manufacturer",
+        manufacturer="manufacturer_orig_1",
-        model="model",
+        model="model_orig_1",
+        model_id="model_id_orig_1",
+        name="name_orig_1",
+        serial_number="serial_no_orig_1",
+        suggested_area="suggested_area_orig_1",
+        sw_version="version_orig_1",
+        via_device="via_device_id_orig_1",
     )
 
     assert len(device_registry.devices) == 1
     assert len(device_registry.deleted_devices) == 0
 
+    # Add another config entry to the same device
     device_registry.async_get_or_create(
         config_entry_id=config_entry_2.entry_id,
+        configuration_url="http://config_url_orig_2.bla",
         connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        entry_type=None,
+        hw_version="hw_version_orig_2",
         identifiers={("entry_234", "2345")},
-        manufacturer="manufacturer",
+        manufacturer="manufacturer_orig_2",
-        model="model",
+        model="model_orig_2",
+        model_id="model_id_orig_2",
+        name="name_orig_2",
+        serial_number="serial_no_orig_2",
+        suggested_area="suggested_area_orig_2",
+        sw_version="version_orig_2",
+        via_device="via_device_id_orig_2",
     )
 
     assert len(device_registry.devices) == 1
     assert len(device_registry.deleted_devices) == 0
 
+    # Apply user customizations
+    updated_device = device_registry.async_update_device(
+        entry.id,
+        area_id="12345A",
+        disabled_by=dr.DeviceEntryDisabler.USER,
+        labels={"label1", "label2"},
+        name_by_user="Test Friendly Name",
+    )
+
+    # Check device entry before we remove it
+    assert updated_device == dr.DeviceEntry(
+        area_id="12345A",
+        config_entries={config_entry_1.entry_id, config_entry_2.entry_id},
+        config_entries_subentries={
+            config_entry_1.entry_id: {"mock-subentry-id-1-1"},
+            config_entry_2.entry_id: {None},
+        },
+        configuration_url="http://config_url_orig_2.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")},
+        created_at=utcnow(),
+        disabled_by=dr.DeviceEntryDisabler.USER,
+        entry_type=None,
+        hw_version="hw_version_orig_2",
+        id=entry.id,
+        identifiers={("entry_123", "0123"), ("entry_234", "2345")},
+        labels={"label1", "label2"},
+        manufacturer="manufacturer_orig_2",
+        model="model_orig_2",
+        model_id="model_id_orig_2",
+        modified_at=utcnow(),
+        name_by_user="Test Friendly Name",
+        name="name_orig_2",
+        primary_config_entry=config_entry_1.entry_id,
+        serial_number="serial_no_orig_2",
+        suggested_area="suggested_area_orig_2",
+        sw_version="version_orig_2",
+    )
+
     device_registry.async_remove_device(entry.id)
 
     assert len(device_registry.devices) == 0
     assert len(device_registry.deleted_devices) == 1
 
+    # config_entry_1 restores the original device, only the supplied config entry,
+    # config subentry, connections, and identifiers will be restored
     entry2 = device_registry.async_get_or_create(
         config_entry_id=config_entry_1.entry_id,
+        config_subentry_id="mock-subentry-id-1-1",
+        configuration_url="http://config_url_new_1.bla",
         connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        entry_type=dr.DeviceEntryType.SERVICE,
+        hw_version="hw_version_new_1",
         identifiers={("entry_123", "0123")},
-        manufacturer="manufacturer",
+        manufacturer="manufacturer_new_1",
-        model="model",
+        model="model_new_1",
+        model_id="model_id_new_1",
+        name="name_new_1",
+        serial_number="serial_no_new_1",
+        suggested_area="suggested_area_new_1",
+        sw_version="version_new_1",
+        via_device="via_device_id_new_1",
+    )
+
+    assert entry2 == dr.DeviceEntry(
+        area_id="suggested_area_new_1",
+        config_entries={config_entry_1.entry_id},
+        config_entries_subentries={config_entry_1.entry_id: {"mock-subentry-id-1-1"}},
+        configuration_url="http://config_url_new_1.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")},
+        created_at=utcnow(),
+        disabled_by=None,
+        entry_type=dr.DeviceEntryType.SERVICE,
+        hw_version="hw_version_new_1",
+        id=entry.id,
+        identifiers={("entry_123", "0123")},
+        labels={},
+        manufacturer="manufacturer_new_1",
+        model="model_new_1",
+        model_id="model_id_new_1",
+        modified_at=utcnow(),
+        name_by_user=None,
+        name="name_new_1",
+        primary_config_entry=config_entry_1.entry_id,
+        serial_number="serial_no_new_1",
+        suggested_area="suggested_area_new_1",
+        sw_version="version_new_1",
     )
 
-    assert entry.id == entry2.id
     assert len(device_registry.devices) == 1
     assert len(device_registry.deleted_devices) == 0
 
@@ -3352,17 +3517,55 @@ async def test_restore_shared_device(
     assert isinstance(entry2.connections, set)
     assert isinstance(entry2.identifiers, set)
 
+    # Remove the device again
     device_registry.async_remove_device(entry.id)
 
+    # config_entry_2 restores the original device, only the supplied config entry,
+    # config subentry, connections, and identifiers will be restored
     entry3 = device_registry.async_get_or_create(
         config_entry_id=config_entry_2.entry_id,
+        configuration_url="http://config_url_new_2.bla",
         connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        entry_type=None,
+        hw_version="hw_version_new_2",
         identifiers={("entry_234", "2345")},
-        manufacturer="manufacturer",
+        manufacturer="manufacturer_new_2",
-        model="model",
+        model="model_new_2",
+        model_id="model_id_new_2",
+        name="name_new_2",
+        serial_number="serial_no_new_2",
+        suggested_area="suggested_area_new_2",
+        sw_version="version_new_2",
+        via_device="via_device_id_new_2",
+    )
+
+    assert entry3 == dr.DeviceEntry(
+        area_id="suggested_area_new_2",
+        config_entries={config_entry_2.entry_id},
+        config_entries_subentries={
+            config_entry_2.entry_id: {None},
+        },
+        configuration_url="http://config_url_new_2.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")},
+        created_at=utcnow(),
+        disabled_by=None,
+        entry_type=None,
+        hw_version="hw_version_new_2",
+        id=entry.id,
+        identifiers={("entry_234", "2345")},
+        labels={},
+        manufacturer="manufacturer_new_2",
+        model="model_new_2",
+        model_id="model_id_new_2",
+        modified_at=utcnow(),
+        name_by_user=None,
+        name="name_new_2",
+        primary_config_entry=config_entry_2.entry_id,
+        serial_number="serial_no_new_2",
+        suggested_area="suggested_area_new_2",
+        sw_version="version_new_2",
     )
 
-    assert entry.id == entry3.id
     assert len(device_registry.devices) == 1
     assert len(device_registry.deleted_devices) == 0
 
@@ -3370,15 +3573,53 @@ async def test_restore_shared_device(
     assert isinstance(entry3.connections, set)
     assert isinstance(entry3.identifiers, set)
 
+    # Add config_entry_1 back to the restored device
     entry4 = device_registry.async_get_or_create(
         config_entry_id=config_entry_1.entry_id,
+        config_subentry_id="mock-subentry-id-1-1",
+        configuration_url="http://config_url_new_1.bla",
         connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        entry_type=dr.DeviceEntryType.SERVICE,
+        hw_version="hw_version_new_1",
         identifiers={("entry_123", "0123")},
-        manufacturer="manufacturer",
+        manufacturer="manufacturer_new_1",
-        model="model",
+        model="model_new_1",
+        model_id="model_id_new_1",
+        name="name_new_1",
+        serial_number="serial_no_new_1",
+        suggested_area="suggested_area_new_1",
+        sw_version="version_new_1",
+        via_device="via_device_id_new_1",
+    )
+
+    assert entry4 == dr.DeviceEntry(
+        area_id="suggested_area_new_2",
+        config_entries={config_entry_1.entry_id, config_entry_2.entry_id},
+        config_entries_subentries={
+            config_entry_1.entry_id: {"mock-subentry-id-1-1"},
+            config_entry_2.entry_id: {None},
+        },
+        configuration_url="http://config_url_new_1.bla",
+        connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:ab:cd:ef")},
+        created_at=utcnow(),
+        disabled_by=None,
+        entry_type=dr.DeviceEntryType.SERVICE,
+        hw_version="hw_version_new_1",
+        id=entry.id,
+        identifiers={("entry_123", "0123"), ("entry_234", "2345")},
+        labels={},
+        manufacturer="manufacturer_new_1",
+        model="model_new_1",
+        model_id="model_id_new_1",
+        modified_at=utcnow(),
+        name_by_user=None,
+        name="name_new_1",
+        primary_config_entry=config_entry_2.entry_id,
+        serial_number="serial_no_new_1",
+        suggested_area="suggested_area_new_1",
+        sw_version="version_new_1",
     )
 
-    assert entry.id == entry4.id
     assert len(device_registry.devices) == 1
     assert len(device_registry.deleted_devices) == 0
 
@@ -3388,7 +3629,7 @@ async def test_restore_shared_device(
 
     await hass.async_block_till_done()
 
-    assert len(update_events) == 7
+    assert len(update_events) == 8
     assert update_events[0].data == {
         "action": "create",
         "device_id": entry.id,
@@ -3398,33 +3639,65 @@ async def test_restore_shared_device(
         "device_id": entry.id,
         "changes": {
             "config_entries": {config_entry_1.entry_id},
-            "config_entries_subentries": {config_entry_1.entry_id: {None}},
+            "config_entries_subentries": {
+                config_entry_1.entry_id: {"mock-subentry-id-1-1"}
+            },
+            "configuration_url": "http://config_url_orig_1.bla",
+            "entry_type": dr.DeviceEntryType.SERVICE,
+            "hw_version": "hw_version_orig_1",
             "identifiers": {("entry_123", "0123")},
+            "manufacturer": "manufacturer_orig_1",
+            "model": "model_orig_1",
+            "model_id": "model_id_orig_1",
+            "name": "name_orig_1",
+            "serial_number": "serial_no_orig_1",
+            "suggested_area": "suggested_area_orig_1",
+            "sw_version": "version_orig_1",
         },
     }
     assert update_events[2].data == {
-        "action": "remove",
+        "action": "update",
         "device_id": entry.id,
+        "changes": {
+            "area_id": "suggested_area_orig_1",
+            "disabled_by": None,
+            "labels": set(),
+            "name_by_user": None,
+        },
     }
     assert update_events[3].data == {
-        "action": "create",
+        "action": "remove",
         "device_id": entry.id,
     }
     assert update_events[4].data == {
-        "action": "remove",
-        "device_id": entry.id,
-    }
-    assert update_events[5].data == {
         "action": "create",
         "device_id": entry.id,
     }
+    assert update_events[5].data == {
+        "action": "remove",
+        "device_id": entry.id,
+    }
     assert update_events[6].data == {
+        "action": "create",
+        "device_id": entry.id,
+    }
+    assert update_events[7].data == {
         "action": "update",
         "device_id": entry.id,
         "changes": {
             "config_entries": {config_entry_2.entry_id},
             "config_entries_subentries": {config_entry_2.entry_id: {None}},
+            "configuration_url": "http://config_url_new_2.bla",
+            "entry_type": None,
+            "hw_version": "hw_version_new_2",
             "identifiers": {("entry_234", "2345")},
+            "manufacturer": "manufacturer_new_2",
+            "model": "model_new_2",
+            "model_id": "model_id_new_2",
+            "name": "name_new_2",
+            "serial_number": "serial_no_new_2",
+            "suggested_area": "suggested_area_new_2",
+            "sw_version": "version_new_2",
         },
     }