Stable device id when a deleted device is restored (#36309)

* Stable device id when a deleted device is restored.

* Tweak

* Store only basic data for deleted devices

* Simplify code

* Simplify code

* Address review comments.

* Improve test

* Fix missing save

homeassistant/helpers/device_registry.py

@@ -33,6 +33,26 @@ CONNECTION_UPNP = "upnp"
 CONNECTION_ZIGBEE = "zigbee"
 
 
+@attr.s(slots=True, frozen=True)
+class DeletedDeviceEntry:
+    """Deleted Device Registry Entry."""
+
+    config_entries: Set[str] = attr.ib()
+    connections: Set[Tuple[str, str]] = attr.ib()
+    identifiers: Set[Tuple[str, str]] = attr.ib()
+    id: str = attr.ib()
+
+    def to_device_entry(self):
+        """Create DeviceEntry from DeletedDeviceEntry."""
+        return DeviceEntry(
+            config_entries=self.config_entries,
+            connections=self.connections,
+            identifiers=self.identifiers,
+            id=self.id,
+            is_new=True,
+        )
+
+
 @attr.s(slots=True, frozen=True)
 class DeviceEntry:
     """Device Registry Entry."""
@@ -81,6 +101,7 @@ class DeviceRegistry:
     """Class to hold a registry of devices."""
 
     devices: Dict[str, DeviceEntry]
+    deleted_devices: Dict[str, DeletedDeviceEntry]
 
     def __init__(self, hass: HomeAssistantType) -> None:
         """Initialize the device registry."""
@@ -104,6 +125,18 @@ class DeviceRegistry:
                 return device
         return None
 
+    @callback
+    def _async_get_deleted_device(
+        self, identifiers: set, connections: set
+    ) -> Optional[DeletedDeviceEntry]:
+        """Check if device has previously been registered."""
+        for device in self.deleted_devices.values():
+            if any(iden in device.identifiers for iden in identifiers) or any(
+                conn in device.connections for conn in connections
+            ):
+                return device
+        return None
+
     @callback
     def async_get_or_create(
         self,
@@ -136,7 +169,12 @@ class DeviceRegistry:
         device = self.async_get_device(identifiers, connections)
 
         if device is None:
-            device = DeviceEntry(is_new=True)
+            deleted_device = self._async_get_deleted_device(identifiers, connections)
+            if deleted_device is None:
+                device = DeviceEntry(is_new=True)
+            else:
+                self.deleted_devices.pop(deleted_device.id)
+                device = deleted_device.to_device_entry()
             self.devices[device.id] = device
 
         if via_device is not None:
@@ -283,7 +321,13 @@ class DeviceRegistry:
     @callback
     def async_remove_device(self, device_id: str) -> None:
         """Remove a device from the device registry."""
-        del self.devices[device_id]
+        device = self.devices.pop(device_id)
+        self.deleted_devices[device_id] = DeletedDeviceEntry(
+            config_entries=device.config_entries,
+            connections=device.connections,
+            identifiers=device.identifiers,
+            id=device.id,
+        )
         self.hass.bus.async_fire(
             EVENT_DEVICE_REGISTRY_UPDATED, {"action": "remove", "device_id": device_id}
         )
@@ -296,6 +340,7 @@ class DeviceRegistry:
         data = await self._store.async_load()
 
         devices = OrderedDict()
+        deleted_devices = OrderedDict()
 
         if data is not None:
             for device in data["devices"]:
@@ -319,8 +364,17 @@ class DeviceRegistry:
                     area_id=device.get("area_id"),
                     name_by_user=device.get("name_by_user"),
                 )
+            # Introduced in 0.111
+            for device in data.get("deleted_devices", []):
+                deleted_devices[device["id"]] = DeletedDeviceEntry(
+                    config_entries=set(device["config_entries"]),
+                    connections={tuple(conn) for conn in device["connections"]},
+                    identifiers={tuple(iden) for iden in device["identifiers"]},
+                    id=device["id"],
+                )
 
         self.devices = devices
+        self.deleted_devices = deleted_devices
 
     @callback
     def async_schedule_save(self) -> None:
@@ -349,6 +403,15 @@ class DeviceRegistry:
             }
             for entry in self.devices.values()
         ]
+        data["deleted_devices"] = [
+            {
+                "config_entries": list(entry.config_entries),
+                "connections": list(entry.connections),
+                "identifiers": list(entry.identifiers),
+                "id": entry.id,
+            }
+            for entry in self.deleted_devices.values()
+        ]
 
         return data
 
@@ -357,6 +420,19 @@ class DeviceRegistry:
         """Clear config entry from registry entries."""
         for device in list(self.devices.values()):
             self._async_update_device(device.id, remove_config_entry_id=config_entry_id)
+        for deleted_device in list(self.deleted_devices.values()):
+            config_entries = deleted_device.config_entries
+            if config_entry_id not in config_entries:
+                continue
+            if config_entries == {config_entry_id}:
+                # Permanently remove the device from the device registry.
+                del self.deleted_devices[deleted_device.id]
+            else:
+                config_entries = config_entries - {config_entry_id}
+                self.deleted_devices[deleted_device.id] = attr.evolve(
+                    deleted_device, config_entries=config_entries
+                )
+            self.async_schedule_save()
 
     @callback
     def async_clear_area_id(self, area_id: str) -> None:

tests/common.py

@@ -381,10 +381,11 @@ def mock_area_registry(hass, mock_entries=None):
     return registry
 
 
-def mock_device_registry(hass, mock_entries=None):
+def mock_device_registry(hass, mock_entries=None, mock_deleted_entries=None):
     """Mock the Device Registry."""
     registry = device_registry.DeviceRegistry(hass)
     registry.devices = mock_entries or OrderedDict()
+    registry.deleted_devices = mock_deleted_entries or OrderedDict()
 
     hass.data[device_registry.DATA_REGISTRY] = registry
     return registry

tests/helpers/test_device_registry.py

@@ -153,11 +153,21 @@ async def test_loading_from_storage(hass, hass_storage):
                     "area_id": "12345A",
                     "name_by_user": "Test Friendly Name",
                 }
-            ]
+            ],
+            "deleted_devices": [
+                {
+                    "config_entries": ["1234"],
+                    "connections": [["Zigbee", "23.45.67.89.01"]],
+                    "id": "bcdefghijklmn",
+                    "identifiers": [["serial", "34:56:AB:CD:EF:12"]],
+                }
+            ],
         },
     }
 
     registry = await device_registry.async_get_registry(hass)
+    assert len(registry.devices) == 1
+    assert len(registry.deleted_devices) == 1
 
     entry = registry.async_get_or_create(
         config_entry_id="1234",
@@ -171,6 +181,20 @@ async def test_loading_from_storage(hass, hass_storage):
     assert entry.name_by_user == "Test Friendly Name"
     assert entry.entry_type == "service"
     assert isinstance(entry.config_entries, set)
+    assert isinstance(entry.connections, set)
+    assert isinstance(entry.identifiers, set)
+
+    entry = registry.async_get_or_create(
+        config_entry_id="1234",
+        connections={("Zigbee", "23.45.67.89.01")},
+        identifiers={("serial", "34:56:AB:CD:EF:12")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+    assert entry.id == "bcdefghijklmn"
+    assert isinstance(entry.config_entries, set)
+    assert isinstance(entry.connections, set)
+    assert isinstance(entry.identifiers, set)
 
 
 async def test_removing_config_entries(hass, registry, update_events):
@@ -224,6 +248,79 @@ async def test_removing_config_entries(hass, registry, update_events):
     assert update_events[4]["device_id"] == entry3.id
 
 
+async def test_deleted_device_removing_config_entries(hass, registry, update_events):
+    """Make sure we do not get duplicate entries."""
+    entry = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+    entry2 = registry.async_get_or_create(
+        config_entry_id="456",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+    entry3 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "34:56:78:CD:EF:12")},
+        identifiers={("bridgeid", "4567")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+
+    assert len(registry.devices) == 2
+    assert len(registry.deleted_devices) == 0
+    assert entry.id == entry2.id
+    assert entry.id != entry3.id
+    assert entry2.config_entries == {"123", "456"}
+
+    registry.async_remove_device(entry.id)
+    registry.async_remove_device(entry3.id)
+
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 2
+
+    await hass.async_block_till_done()
+    assert len(update_events) == 5
+    assert update_events[0]["action"] == "create"
+    assert update_events[0]["device_id"] == entry.id
+    assert update_events[1]["action"] == "update"
+    assert update_events[1]["device_id"] == entry2.id
+    assert update_events[2]["action"] == "create"
+    assert update_events[2]["device_id"] == entry3.id
+    assert update_events[3]["action"] == "remove"
+    assert update_events[3]["device_id"] == entry.id
+    assert update_events[4]["action"] == "remove"
+    assert update_events[4]["device_id"] == entry3.id
+
+    registry.async_clear_config_entry("123")
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 1
+
+    registry.async_clear_config_entry("456")
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 0
+
+    # No event when a deleted device is purged
+    await hass.async_block_till_done()
+    assert len(update_events) == 5
+
+    # Re-add, expect new device id
+    entry2 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+
+    assert entry.id != entry2.id
+
+
 async def test_removing_area_id(registry):
     """Make sure we can clear area id."""
     entry = registry.async_get_or_create(
@@ -243,6 +340,36 @@ async def test_removing_area_id(registry):
     assert entry_w_area != entry_wo_area
 
 
+async def test_deleted_device_removing_area_id(registry):
+    """Make sure we can clear area id of deleted device."""
+    entry = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+
+    entry_w_area = registry.async_update_device(entry.id, area_id="12345A")
+
+    registry.async_remove_device(entry.id)
+    registry.async_clear_area_id("12345A")
+
+    entry2 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+    assert entry.id == entry2.id
+
+    entry_wo_area = registry.async_get_device({("bridgeid", "0123")}, set())
+
+    assert not entry_wo_area.area_id
+    assert entry_w_area != entry_wo_area
+
+
 async def test_specifying_via_device_create(registry):
     """Test specifying a via_device and updating."""
     via = registry.async_get_or_create(
@@ -320,7 +447,19 @@ async def test_loading_saving_data(hass, registry):
         via_device=("hue", "0123"),
     )
 
+    orig_light2 = registry.async_get_or_create(
+        config_entry_id="456",
+        connections=set(),
+        identifiers={("hue", "789")},
+        manufacturer="manufacturer",
+        model="light",
+        via_device=("hue", "0123"),
+    )
+
+    registry.async_remove_device(orig_light2.id)
+
     assert len(registry.devices) == 2
+    assert len(registry.deleted_devices) == 1
 
     orig_via = registry.async_update_device(
         orig_via.id, area_id="mock-area-id", name_by_user="mock-name-by-user"
@@ -333,6 +472,7 @@ async def test_loading_saving_data(hass, registry):
 
     # Ensure same order
     assert list(registry.devices) == list(registry2.devices)
+    assert list(registry.deleted_devices) == list(registry2.deleted_devices)
 
     new_via = registry2.async_get_device({("hue", "0123")}, set())
     new_light = registry2.async_get_device({("hue", "456")}, set())
@@ -584,3 +724,103 @@ async def test_cleanup_entity_registry_change(hass):
         ent_reg.async_remove(entity.entity_id)
         await hass.async_block_till_done()
         assert len(mock_call.mock_calls) == 2
+
+
+async def test_restore_device(hass, registry, update_events):
+    """Make sure device id is stable."""
+    entry = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+
+    assert len(registry.devices) == 1
+    assert len(registry.deleted_devices) == 0
+
+    registry.async_remove_device(entry.id)
+
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 1
+
+    entry2 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "34:56:78:CD:EF:12")},
+        identifiers={("bridgeid", "4567")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+    entry3 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+        manufacturer="manufacturer",
+        model="model",
+    )
+
+    assert entry.id == entry3.id
+    assert entry.id != entry2.id
+    assert len(registry.devices) == 2
+    assert len(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) == 4
+    assert update_events[0]["action"] == "create"
+    assert update_events[0]["device_id"] == entry.id
+    assert update_events[1]["action"] == "remove"
+    assert update_events[1]["device_id"] == entry.id
+    assert update_events[2]["action"] == "create"
+    assert update_events[2]["device_id"] == entry2.id
+    assert update_events[3]["action"] == "create"
+    assert update_events[3]["device_id"] == entry3.id
+
+
+async def test_restore_simple_device(hass, registry, update_events):
+    """Make sure device id is stable."""
+    entry = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+    )
+
+    assert len(registry.devices) == 1
+    assert len(registry.deleted_devices) == 0
+
+    registry.async_remove_device(entry.id)
+
+    assert len(registry.devices) == 0
+    assert len(registry.deleted_devices) == 1
+
+    entry2 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "34:56:78:CD:EF:12")},
+        identifiers={("bridgeid", "4567")},
+    )
+    entry3 = registry.async_get_or_create(
+        config_entry_id="123",
+        connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
+        identifiers={("bridgeid", "0123")},
+    )
+
+    assert entry.id == entry3.id
+    assert entry.id != entry2.id
+    assert len(registry.devices) == 2
+    assert len(registry.deleted_devices) == 0
+
+    await hass.async_block_till_done()
+
+    assert len(update_events) == 4
+    assert update_events[0]["action"] == "create"
+    assert update_events[0]["device_id"] == entry.id
+    assert update_events[1]["action"] == "remove"
+    assert update_events[1]["device_id"] == entry.id
+    assert update_events[2]["action"] == "create"
+    assert update_events[2]["device_id"] == entry2.id
+    assert update_events[3]["action"] == "create"
+    assert update_events[3]["device_id"] == entry3.id