From 8f04f22c65bcea7c2359229693232415e85bc48d Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Fri, 29 Aug 2025 16:48:29 +0200 Subject: [PATCH] Improve migration to device registry version 1.11 (#151315) Co-authored-by: Franck Nijhof --- homeassistant/helpers/device_registry.py | 49 ++++++--- tests/helpers/test_device_registry.py | 131 ++++++++++++++++++++++- 2 files changed, 165 insertions(+), 15 deletions(-) diff --git a/homeassistant/helpers/device_registry.py b/homeassistant/helpers/device_registry.py index fd11f7b5f21..aa619c1dc41 100644 --- a/homeassistant/helpers/device_registry.py +++ b/homeassistant/helpers/device_registry.py @@ -9,7 +9,7 @@ from enum import StrEnum from functools import lru_cache import logging import time -from typing import TYPE_CHECKING, Any, Literal, TypedDict +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict import attr from yarl import URL @@ -68,6 +68,8 @@ CONNECTION_ZIGBEE = "zigbee" ORPHANED_DEVICE_KEEP_SECONDS = 86400 * 30 +UNDEFINED_STR: Final = "UNDEFINED" + # Can be removed when suggested_area is removed from DeviceEntry RUNTIME_ONLY_ATTRS = {"suggested_area"} @@ -463,7 +465,7 @@ class DeletedDeviceEntry: validator=_normalize_connections_validator ) created_at: datetime = attr.ib() - disabled_by: DeviceEntryDisabler | None = attr.ib() + disabled_by: DeviceEntryDisabler | UndefinedType | None = attr.ib() id: str = attr.ib() identifiers: set[tuple[str, str]] = attr.ib() labels: set[str] = attr.ib() @@ -478,15 +480,19 @@ class DeletedDeviceEntry: config_subentry_id: str | None, connections: set[tuple[str, str]], identifiers: set[tuple[str, str]], + disabled_by: DeviceEntryDisabler | UndefinedType | None, ) -> 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 + if self.disabled_by is not UNDEFINED: + 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 + else: + disabled_by = disabled_by if disabled_by is not UNDEFINED else None return DeviceEntry( area_id=self.area_id, # type ignores: likely https://github.com/python/mypy/issues/8625 @@ -517,7 +523,9 @@ class DeletedDeviceEntry: }, "connections": list(self.connections), "created_at": self.created_at, - "disabled_by": self.disabled_by, + "disabled_by": self.disabled_by + if self.disabled_by is not UNDEFINED + else UNDEFINED_STR, "identifiers": list(self.identifiers), "id": self.id, "labels": list(self.labels), @@ -605,7 +613,7 @@ class DeviceRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]): # Introduced in 2025.6 for device in old_data["deleted_devices"]: device["area_id"] = None - device["disabled_by"] = None + device["disabled_by"] = UNDEFINED_STR device["labels"] = [] device["name_by_user"] = None if old_minor_version < 11: @@ -935,6 +943,7 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]): config_subentry_id if config_subentry_id is not UNDEFINED else None, connections, identifiers, + disabled_by, ) disabled_by = UNDEFINED @@ -1502,7 +1511,21 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]): sw_version=device["sw_version"], via_device_id=device["via_device_id"], ) + # Introduced in 0.111 + def get_optional_enum[_EnumT: StrEnum]( + cls: type[_EnumT], value: str | None + ) -> _EnumT | UndefinedType | None: + """Convert string to the passed enum, UNDEFINED or None.""" + if value is None: + return None + if value == UNDEFINED_STR: + return UNDEFINED + try: + return cls(value) + except ValueError: + return None + for device in data["deleted_devices"]: deleted_devices[device["id"]] = DeletedDeviceEntry( area_id=device["area_id"], @@ -1515,10 +1538,8 @@ class DeviceRegistry(BaseRegistry[dict[str, list[dict[str, Any]]]]): }, connections={tuple(conn) for conn in device["connections"]}, created_at=datetime.fromisoformat(device["created_at"]), - disabled_by=( - DeviceEntryDisabler(device["disabled_by"]) - if device["disabled_by"] - else None + disabled_by=get_optional_enum( + DeviceEntryDisabler, device["disabled_by"] ), identifiers={tuple(iden) for iden in device["identifiers"]}, id=device["id"], diff --git a/tests/helpers/test_device_registry.py b/tests/helpers/test_device_registry.py index 9690b2a52fa..8cfd3c66ad9 100644 --- a/tests/helpers/test_device_registry.py +++ b/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 @@ -508,6 +510,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, @@ -581,7 +586,7 @@ async def test_migration_from_1_1( "config_entries_subentries": {"123456": [None]}, "connections": [], "created_at": "1970-01-01T00:00:00+00:00", - "disabled_by": None, + "disabled_by": "UNDEFINED", "id": "deletedid", "identifiers": [["serial", "123456ABCDFF"]], "labels": [], @@ -3833,6 +3838,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",