Adjust device disabled_by flag when restoring a deleted device (#151154)

homeassistant/helpers/device_registry.py

@@ -476,20 +476,27 @@ class DeletedDeviceEntry:
 
     def to_device_entry(
         self,
-        config_entry_id: str,
+        config_entry: ConfigEntry,
         config_subentry_id: str | None,
         connections: set[tuple[str, str]],
         identifiers: set[tuple[str, str]],
     ) -> DeviceEntry:
         """Create DeviceEntry from DeletedDeviceEntry."""
+        # Adjust disabled_by based on config entry state
+        disabled_by = self.disabled_by
+        if config_entry.disabled_by:
+            if disabled_by is None:
+                disabled_by = DeviceEntryDisabler.CONFIG_ENTRY
+        elif disabled_by == DeviceEntryDisabler.CONFIG_ENTRY:
+            disabled_by = None
         return DeviceEntry(
             area_id=self.area_id,
             # type ignores: likely https://github.com/python/mypy/issues/8625
-            config_entries={config_entry_id},  # type: ignore[arg-type]
-            config_entries_subentries={config_entry_id: {config_subentry_id}},
+            config_entries={config_entry.entry_id},  # type: ignore[arg-type]
+            config_entries_subentries={config_entry.entry_id: {config_subentry_id}},
             connections=self.connections & connections,  # type: ignore[arg-type]
             created_at=self.created_at,
-            disabled_by=self.disabled_by,
+            disabled_by=disabled_by,
             identifiers=self.identifiers & identifiers,  # type: ignore[arg-type]
             id=self.id,
             is_new=True,
@@ -922,7 +929,7 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]):
             else:
                 self.deleted_devices.pop(deleted_device.id)
                 device = deleted_device.to_device_entry(
-                    config_entry_id,
+                    config_entry,
                     # Interpret not specifying a subentry as None
                     config_subentry_id if config_subentry_id is not UNDEFINED else None,
                     connections,

tests/helpers/test_device_registry.py

@@ -3573,6 +3573,173 @@ async def test_restore_device(
     }
 
 
+@pytest.mark.parametrize(
+    (
+        "config_entry_disabled_by",
+        "device_disabled_by_initial",
+        "device_disabled_by_restored",
+    ),
+    [
+        (
+            None,
+            None,
+            None,
+        ),
+        # Config entry not disabled, device was disabled by config entry.
+        # Device not disabled when restored.
+        (
+            None,
+            dr.DeviceEntryDisabler.CONFIG_ENTRY,
+            None,
+        ),
+        (
+            None,
+            dr.DeviceEntryDisabler.INTEGRATION,
+            dr.DeviceEntryDisabler.INTEGRATION,
+        ),
+        (
+            None,
+            dr.DeviceEntryDisabler.USER,
+            dr.DeviceEntryDisabler.USER,
+        ),
+        # Config entry disabled, device not disabled.
+        # Device disabled by config entry when restored.
+        (
+            config_entries.ConfigEntryDisabler.USER,
+            None,
+            dr.DeviceEntryDisabler.CONFIG_ENTRY,
+        ),
+        (
+            config_entries.ConfigEntryDisabler.USER,
+            dr.DeviceEntryDisabler.CONFIG_ENTRY,
+            dr.DeviceEntryDisabler.CONFIG_ENTRY,
+        ),
+        (
+            config_entries.ConfigEntryDisabler.USER,
+            dr.DeviceEntryDisabler.INTEGRATION,
+            dr.DeviceEntryDisabler.INTEGRATION,
+        ),
+        (
+            config_entries.ConfigEntryDisabler.USER,
+            dr.DeviceEntryDisabler.USER,
+            dr.DeviceEntryDisabler.USER,
+        ),
+    ],
+)
+@pytest.mark.usefixtures("freezer")
+async def test_restore_disabled_by(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    mock_config_entry: MockConfigEntry,
+    config_entry_disabled_by: config_entries.ConfigEntryDisabler | None,
+    device_disabled_by_initial: dr.DeviceEntryDisabler | None,
+    device_disabled_by_restored: 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)
+    await hass.config_entries.async_set_disabled_by(
+        mock_config_entry.entry_id, config_entry_disabled_by
+    )
+    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=device_disabled_by_initial,
+        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 entry.disabled_by == device_disabled_by_initial
+
+    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
+
+    # 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")},
+        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=device_disabled_by_restored,
+        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.usefixtures("freezer")
 async def test_restore_shared_device(
     hass: HomeAssistant, device_registry: dr.DeviceRegistry