Improve migration to device registry version 1.11 (#151315)

Co-authored-by: Franck Nijhof <git@frenck.dev>
This commit is contained in:
Erik Montnemery
2025-08-29 16:48:29 +02:00
committed by GitHub
parent a01f638fc6
commit 8f04f22c65
2 changed files with 165 additions and 15 deletions

View File

@@ -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"],

View File

@@ -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",