From 50108e23ed53f745382997221e48799403824b6a Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Mon, 25 Aug 2025 22:49:47 +0200 Subject: [PATCH] Adjust device disabled_by flag when restoring a deleted device (#151154) --- homeassistant/helpers/device_registry.py | 17 ++- tests/helpers/test_device_registry.py | 167 +++++++++++++++++++++++ 2 files changed, 179 insertions(+), 5 deletions(-) diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index c7f7d4c369d..a78fa935606 100644 --- a/homeassistant/helpers/device_registry.py +++ b/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, diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py index d056c25fc3b..30fb22b8e09 100644 --- a/tests/helpers/test_device_registry.py +++ b/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