Add config subentry support to entity registry

pull/132049/head
Erik 2024-10-11 12:53:02 +02:00
parent 90b2504d5a
commit 68f8c3e9ed
3 changed files with 729 additions and 10 deletions
homeassistant/helpers

View File

@ -79,7 +79,7 @@ EVENT_ENTITY_REGISTRY_UPDATED: EventType[EventEntityRegistryUpdatedData] = Event
_LOGGER = logging.getLogger(__name__)
STORAGE_VERSION_MAJOR = 1
STORAGE_VERSION_MINOR = 15
STORAGE_VERSION_MINOR = 16
STORAGE_KEY = "core.entity_registry"
CLEANUP_INTERVAL = 3600 * 24
@ -179,6 +179,7 @@ class RegistryEntry:
categories: dict[str, str] = attr.ib(factory=dict)
capabilities: Mapping[str, Any] | None = attr.ib(default=None)
config_entry_id: str | None = attr.ib(default=None)
config_subentry_id: str | None = attr.ib(default=None)
created_at: datetime = attr.ib(factory=utcnow)
device_class: str | None = attr.ib(default=None)
device_id: str | None = attr.ib(default=None)
@ -282,6 +283,7 @@ class RegistryEntry:
"area_id": self.area_id,
"categories": self.categories,
"config_entry_id": self.config_entry_id,
"config_subentry_id": self.config_subentry_id,
"created_at": self.created_at.timestamp(),
"device_id": self.device_id,
"disabled_by": self.disabled_by,
@ -343,6 +345,7 @@ class RegistryEntry:
"categories": self.categories,
"capabilities": self.capabilities,
"config_entry_id": self.config_entry_id,
"config_subentry_id": self.config_subentry_id,
"created_at": self.created_at,
"device_class": self.device_class,
"device_id": self.device_id,
@ -407,6 +410,7 @@ class DeletedRegistryEntry:
unique_id: str = attr.ib()
platform: str = attr.ib()
config_entry_id: str | None = attr.ib()
config_subentry_id: str | None = attr.ib()
domain: str = attr.ib(init=False, repr=False)
id: str = attr.ib()
orphaned_timestamp: float | None = attr.ib()
@ -426,6 +430,7 @@ class DeletedRegistryEntry:
json_bytes(
{
"config_entry_id": self.config_entry_id,
"config_subentry_id": self.config_subentry_id,
"created_at": self.created_at,
"entity_id": self.entity_id,
"id": self.id,
@ -541,6 +546,13 @@ class EntityRegistryStore(storage.Store[dict[str, list[dict[str, Any]]]]):
for entity in data["deleted_entities"]:
entity["created_at"] = entity["modified_at"] = created_at
if old_minor_version < 16:
# Version 1.16 adds config_subentry_id
for entity in data["entities"]:
entity["config_subentry_id"] = None
for entity in data["deleted_entities"]:
entity["config_subentry_id"] = None
if old_major_version > 1:
raise NotImplementedError
return data
@ -648,9 +660,12 @@ def _validate_item(
domain: str,
platform: str,
*,
config_entry_id: str | None | UndefinedType = None,
config_subentry_id: str | None | UndefinedType = None,
disabled_by: RegistryEntryDisabler | None | UndefinedType = None,
entity_category: EntityCategory | None | UndefinedType = None,
hidden_by: RegistryEntryHider | None | UndefinedType = None,
old_config_subentry_id: str | None = None,
report_non_string_unique_id: bool = True,
unique_id: str | Hashable | UndefinedType | Any,
) -> None:
@ -671,6 +686,26 @@ def _validate_item(
unique_id,
report_issue,
)
if (
config_entry_id
and config_entry_id is not UNDEFINED
and old_config_subentry_id
and config_subentry_id is UNDEFINED
):
raise ValueError("Can't change config entry without changing subentry")
if (
config_entry_id
and config_entry_id is not UNDEFINED
and config_subentry_id
and config_subentry_id is not UNDEFINED
):
if (
not (config_entry := hass.config_entries.async_get_entry(config_entry_id))
or config_subentry_id not in config_entry.subentries
):
raise ValueError(
f"Config entry {config_entry_id} has no subentry {config_subentry_id}"
)
if (
disabled_by
and disabled_by is not UNDEFINED
@ -817,6 +852,7 @@ class EntityRegistry(BaseRegistry):
# Data that we want entry to have
capabilities: Mapping[str, Any] | None | UndefinedType = UNDEFINED,
config_entry: ConfigEntry | None | UndefinedType = UNDEFINED,
config_subentry_id: str | None | UndefinedType = UNDEFINED,
device_id: str | None | UndefinedType = UNDEFINED,
entity_category: EntityCategory | UndefinedType | None = UNDEFINED,
has_entity_name: bool | UndefinedType = UNDEFINED,
@ -843,6 +879,7 @@ class EntityRegistry(BaseRegistry):
entity_id,
capabilities=capabilities,
config_entry_id=config_entry_id,
config_subentry_id=config_subentry_id,
device_id=device_id,
entity_category=entity_category,
has_entity_name=has_entity_name,
@ -859,6 +896,8 @@ class EntityRegistry(BaseRegistry):
self.hass,
domain,
platform,
config_entry_id=config_entry_id,
config_subentry_id=config_subentry_id,
disabled_by=disabled_by,
entity_category=entity_category,
hidden_by=hidden_by,
@ -896,6 +935,7 @@ class EntityRegistry(BaseRegistry):
entry = RegistryEntry(
capabilities=none_if_undefined(capabilities),
config_entry_id=none_if_undefined(config_entry_id),
config_subentry_id=none_if_undefined(config_subentry_id),
created_at=created_at,
device_id=none_if_undefined(device_id),
disabled_by=disabled_by,
@ -938,6 +978,7 @@ class EntityRegistry(BaseRegistry):
orphaned_timestamp = None if config_entry_id else time.time()
self.deleted_entities[key] = DeletedRegistryEntry(
config_entry_id=config_entry_id,
config_subentry_id=entity.config_subentry_id,
created_at=entity.created_at,
entity_id=entity_id,
id=entity.id,
@ -997,6 +1038,19 @@ class EntityRegistry(BaseRegistry):
):
self.async_remove(entity.entity_id)
# Remove entities which belong to config subentries no longer associated with the
# device
entities = async_entries_for_device(
self, event.data["device_id"], include_disabled_entities=True
)
for entity in entities:
if (
(entry_id := entity.config_entry_id) is not None
and entry_id in device.config_entries
and entity.config_subentry_id not in device.config_subentries[entry_id]
):
self.async_remove(entity.entity_id)
# Re-enable disabled entities if the device is no longer disabled
if not device.disabled:
entities = async_entries_for_device(
@ -1030,6 +1084,7 @@ class EntityRegistry(BaseRegistry):
categories: dict[str, str] | UndefinedType = UNDEFINED,
capabilities: Mapping[str, Any] | None | UndefinedType = UNDEFINED,
config_entry_id: str | None | UndefinedType = UNDEFINED,
config_subentry_id: str | None | UndefinedType = UNDEFINED,
device_class: str | None | UndefinedType = UNDEFINED,
device_id: str | None | UndefinedType = UNDEFINED,
disabled_by: RegistryEntryDisabler | None | UndefinedType = UNDEFINED,
@ -1062,6 +1117,7 @@ class EntityRegistry(BaseRegistry):
("categories", categories),
("capabilities", capabilities),
("config_entry_id", config_entry_id),
("config_subentry_id", config_subentry_id),
("device_class", device_class),
("device_id", device_id),
("disabled_by", disabled_by),
@ -1090,9 +1146,12 @@ class EntityRegistry(BaseRegistry):
self.hass,
old.domain,
old.platform,
config_entry_id=config_entry_id,
config_subentry_id=config_subentry_id,
disabled_by=disabled_by,
entity_category=entity_category,
hidden_by=hidden_by,
old_config_subentry_id=old.config_subentry_id,
unique_id=new_unique_id,
)
@ -1157,6 +1216,7 @@ class EntityRegistry(BaseRegistry):
categories: dict[str, str] | UndefinedType = UNDEFINED,
capabilities: Mapping[str, Any] | None | UndefinedType = UNDEFINED,
config_entry_id: str | None | UndefinedType = UNDEFINED,
config_subentry_id: str | None | UndefinedType = UNDEFINED,
device_class: str | None | UndefinedType = UNDEFINED,
device_id: str | None | UndefinedType = UNDEFINED,
disabled_by: RegistryEntryDisabler | None | UndefinedType = UNDEFINED,
@ -1183,6 +1243,7 @@ class EntityRegistry(BaseRegistry):
categories=categories,
capabilities=capabilities,
config_entry_id=config_entry_id,
config_subentry_id=config_subentry_id,
device_class=device_class,
device_id=device_id,
disabled_by=disabled_by,
@ -1209,6 +1270,7 @@ class EntityRegistry(BaseRegistry):
new_platform: str,
*,
new_config_entry_id: str | UndefinedType = UNDEFINED,
new_config_subentry_id: str | UndefinedType = UNDEFINED,
new_unique_id: str | UndefinedType = UNDEFINED,
new_device_id: str | None | UndefinedType = UNDEFINED,
) -> RegistryEntry:
@ -1233,6 +1295,7 @@ class EntityRegistry(BaseRegistry):
entity_id,
new_unique_id=new_unique_id,
config_entry_id=new_config_entry_id,
config_subentry_id=new_config_subentry_id,
device_id=new_device_id,
platform=new_platform,
)
@ -1295,6 +1358,7 @@ class EntityRegistry(BaseRegistry):
categories=entity["categories"],
capabilities=entity["capabilities"],
config_entry_id=entity["config_entry_id"],
config_subentry_id=entity["config_subentry_id"],
created_at=datetime.fromisoformat(entity["created_at"]),
device_class=entity["device_class"],
device_id=entity["device_id"],
@ -1344,6 +1408,7 @@ class EntityRegistry(BaseRegistry):
)
deleted_entities[key] = DeletedRegistryEntry(
config_entry_id=entity["config_entry_id"],
config_subentry_id=entity["config_subentry_id"],
created_at=datetime.fromisoformat(entity["created_at"]),
entity_id=entity["entity_id"],
id=entity["id"],
@ -1402,6 +1467,30 @@ class EntityRegistry(BaseRegistry):
)
self.async_schedule_save()
@callback
def async_clear_config_subentry(
self, config_entry_id: str, config_subentry_id: str
) -> None:
"""Clear config subentry from registry entries."""
now_time = time.time()
for entity_id in [
entry.entity_id
for entry in self.entities.get_entries_for_config_entry_id(config_entry_id)
if entry.config_subentry_id == config_subentry_id
]:
self.async_remove(entity_id)
for key, deleted_entity in list(self.deleted_entities.items()):
if config_subentry_id != deleted_entity.config_subentry_id:
continue
# Add a time stamp when the deleted entity became orphaned
self.deleted_entities[key] = attr.evolve(
deleted_entity,
orphaned_timestamp=now_time,
config_entry_id=None,
config_subentry_id=None,
)
self.async_schedule_save()
@callback
def async_purge_expired_orphaned_entities(self) -> None:
"""Purge expired orphaned entities from the registry.

View File

@ -67,6 +67,7 @@ async def test_list_entities(
"area_id": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": utcnow().timestamp(),
"device_id": None,
"disabled_by": None,
@ -89,6 +90,7 @@ async def test_list_entities(
"area_id": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": utcnow().timestamp(),
"device_id": None,
"disabled_by": None,
@ -138,6 +140,7 @@ async def test_list_entities(
"area_id": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": utcnow().timestamp(),
"device_id": None,
"disabled_by": None,
@ -374,6 +377,7 @@ async def test_get_entity(hass: HomeAssistant, client: MockHAClientWebSocket) ->
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": name_created_at.timestamp(),
"device_class": None,
"device_id": None,
@ -410,6 +414,7 @@ async def test_get_entity(hass: HomeAssistant, client: MockHAClientWebSocket) ->
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": no_name_created_at.timestamp(),
"device_class": None,
"device_id": None,
@ -477,6 +482,7 @@ async def test_get_entities(hass: HomeAssistant, client: MockHAClientWebSocket)
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": name_created_at.timestamp(),
"device_class": None,
"device_id": None,
@ -504,6 +510,7 @@ async def test_get_entities(hass: HomeAssistant, client: MockHAClientWebSocket)
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": no_name_created_at.timestamp(),
"device_class": None,
"device_id": None,
@ -586,6 +593,7 @@ async def test_update_entity(
"categories": {"scope1": "id", "scope2": "id"},
"created_at": created.timestamp(),
"config_entry_id": None,
"config_subentry_id": None,
"device_class": "custom_device_class",
"device_id": None,
"disabled_by": None,
@ -668,6 +676,7 @@ async def test_update_entity(
"capabilities": None,
"categories": {"scope1": "id", "scope2": "id"},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": "custom_device_class",
"device_id": None,
@ -714,6 +723,7 @@ async def test_update_entity(
"capabilities": None,
"categories": {"scope1": "id", "scope2": "id"},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": "custom_device_class",
"device_id": None,
@ -759,6 +769,7 @@ async def test_update_entity(
"capabilities": None,
"categories": {"scope1": "id", "scope2": "id", "scope3": "id"},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": "custom_device_class",
"device_id": None,
@ -804,6 +815,7 @@ async def test_update_entity(
"capabilities": None,
"categories": {"scope1": "id", "scope2": "id", "scope3": "other_id"},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": "custom_device_class",
"device_id": None,
@ -849,6 +861,7 @@ async def test_update_entity(
"capabilities": None,
"categories": {"scope1": "id", "scope3": "other_id"},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": "custom_device_class",
"device_id": None,
@ -911,6 +924,7 @@ async def test_update_entity_require_restart(
"capabilities": None,
"categories": {},
"config_entry_id": config_entry.entry_id,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": None,
"device_id": None,
@ -1032,6 +1046,7 @@ async def test_update_entity_no_changes(
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": None,
"device_id": None,
@ -1129,6 +1144,7 @@ async def test_update_entity_id(
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": created.timestamp(),
"device_class": None,
"device_id": None,

View File

@ -72,11 +72,24 @@ def test_get_or_create_suggested_object_id(entity_registry: er.EntityRegistry) -
def test_get_or_create_updates_data(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
freezer: FrozenDateTimeFactory,
) -> None:
"""Test that we update data in get_or_create."""
orig_config_entry = MockConfigEntry(domain="light")
config_subentry_id = "blabla"
orig_config_entry = MockConfigEntry(
domain="light",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id=config_subentry_id,
title="Mock title",
unique_id="test",
)
],
)
orig_config_entry.add_to_hass(hass)
created = datetime.fromisoformat("2024-02-14T12:00:00.0+00:00")
freezer.move_to(created)
@ -86,6 +99,7 @@ def test_get_or_create_updates_data(
"5678",
capabilities={"max": 100},
config_entry=orig_config_entry,
config_subentry_id=config_subentry_id,
device_id="mock-dev-id",
disabled_by=er.RegistryEntryDisabler.HASS,
entity_category=EntityCategory.CONFIG,
@ -107,6 +121,7 @@ def test_get_or_create_updates_data(
"hue",
capabilities={"max": 100},
config_entry_id=orig_config_entry.entry_id,
config_subentry_id=config_subentry_id,
created_at=created,
device_class=None,
device_id="mock-dev-id",
@ -136,6 +151,7 @@ def test_get_or_create_updates_data(
"5678",
capabilities={"new-max": 150},
config_entry=new_config_entry,
config_subentry_id=None,
device_id="new-mock-dev-id",
disabled_by=er.RegistryEntryDisabler.USER,
entity_category=EntityCategory.DIAGNOSTIC,
@ -157,6 +173,7 @@ def test_get_or_create_updates_data(
area_id=None,
capabilities={"new-max": 150},
config_entry_id=new_config_entry.entry_id,
config_subentry_id=None,
created_at=created,
device_class=None,
device_id="new-mock-dev-id",
@ -476,6 +493,7 @@ async def test_load_bad_data(
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": "2024-02-14T12:00:00.900075+00:00",
"device_class": None,
"device_id": None,
@ -506,6 +524,7 @@ async def test_load_bad_data(
"capabilities": None,
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": "2024-02-14T12:00:00.900075+00:00",
"device_class": None,
"device_id": None,
@ -534,6 +553,7 @@ async def test_load_bad_data(
"deleted_entities": [
{
"config_entry_id": None,
"config_subentry_id": None,
"created_at": "2024-02-14T12:00:00.900075+00:00",
"entity_id": "test.test3",
"id": "00003",
@ -544,6 +564,7 @@ async def test_load_bad_data(
},
{
"config_entry_id": None,
"config_subentry_id": None,
"created_at": "2024-02-14T12:00:00.900075+00:00",
"entity_id": "test.test4",
"id": "00004",
@ -685,6 +706,115 @@ async def test_deleted_entity_removing_config_entry_id(
assert entity_registry.deleted_entities[("light", "hue", "1234")] == deleted_entry2
async def test_removing_config_subentry_id(
hass: HomeAssistant, entity_registry: er.EntityRegistry
) -> None:
"""Test that we update config subentry id in registry."""
update_events = async_capture_events(hass, er.EVENT_ENTITY_REGISTRY_UPDATED)
mock_config = MockConfigEntry(
domain="light",
entry_id="mock-id-1",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
)
],
)
mock_config.add_to_hass(hass)
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=mock_config,
config_subentry_id="mock-subentry-id-1",
)
assert entry.config_subentry_id == "mock-subentry-id-1"
entity_registry.async_clear_config_subentry("mock-id-1", "mock-subentry-id-1")
assert not entity_registry.entities
await hass.async_block_till_done()
assert len(update_events) == 2
assert update_events[0].data == {
"action": "create",
"entity_id": entry.entity_id,
}
assert update_events[1].data == {
"action": "remove",
"entity_id": entry.entity_id,
}
async def test_deleted_entity_removing_config_subentry_id(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
) -> None:
"""Test that we update config subentry id in registry on deleted entity."""
mock_config = MockConfigEntry(
domain="light",
entry_id="mock-id-1",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
),
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-2",
title="Mock title",
unique_id="test",
),
],
)
mock_config.add_to_hass(hass)
entry1 = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=mock_config,
config_subentry_id="mock-subentry-id-1",
)
assert entry1.config_subentry_id == "mock-subentry-id-1"
entry2 = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
config_entry=mock_config,
config_subentry_id="mock-subentry-id-2",
)
assert entry2.config_subentry_id == "mock-subentry-id-2"
entity_registry.async_remove(entry1.entity_id)
entity_registry.async_remove(entry2.entity_id)
assert len(entity_registry.entities) == 0
assert len(entity_registry.deleted_entities) == 2
deleted_entry1 = entity_registry.deleted_entities[("light", "hue", "5678")]
assert deleted_entry1.config_entry_id == "mock-id-1"
assert deleted_entry1.config_subentry_id == "mock-subentry-id-1"
assert deleted_entry1.orphaned_timestamp is None
deleted_entry2 = entity_registry.deleted_entities[("light", "hue", "1234")]
assert deleted_entry2.config_entry_id == "mock-id-1"
assert deleted_entry2.config_subentry_id == "mock-subentry-id-2"
assert deleted_entry2.orphaned_timestamp is None
entity_registry.async_clear_config_subentry("mock-id-1", "mock-subentry-id-1")
assert len(entity_registry.entities) == 0
assert len(entity_registry.deleted_entities) == 2
deleted_entry1 = entity_registry.deleted_entities[("light", "hue", "5678")]
assert deleted_entry1.config_entry_id is None
assert deleted_entry1.config_subentry_id is None
assert deleted_entry1.orphaned_timestamp is not None
assert entity_registry.deleted_entities[("light", "hue", "1234")] == deleted_entry2
async def test_removing_area_id(entity_registry: er.EntityRegistry) -> None:
"""Make sure we can clear area id."""
entry = entity_registry.async_get_or_create("light", "hue", "5678")
@ -740,6 +870,7 @@ async def test_migration_1_1(hass: HomeAssistant, hass_storage: dict[str, Any])
"capabilities": {},
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": "1970-01-01T00:00:00+00:00",
"device_id": None,
"disabled_by": None,
@ -918,6 +1049,7 @@ async def test_migration_1_11(
"capabilities": {},
"categories": {},
"config_entry_id": None,
"config_subentry_id": None,
"created_at": "1970-01-01T00:00:00+00:00",
"device_id": None,
"disabled_by": None,
@ -946,6 +1078,7 @@ async def test_migration_1_11(
"deleted_entities": [
{
"config_entry_id": None,
"config_subentry_id": None,
"created_at": "1970-01-01T00:00:00+00:00",
"entity_id": "test.deleted_entity",
"id": "23456",
@ -1392,7 +1525,7 @@ async def test_remove_config_entry_from_device_removes_entities_2(
config_entry_2.entry_id,
}
# Create one entity for each config entry
# Create an entity without config entry
entry_1 = entity_registry.async_get_or_create(
"light",
"hue",
@ -1412,6 +1545,204 @@ async def test_remove_config_entry_from_device_removes_entities_2(
assert entity_registry.async_is_registered(entry_1.entity_id)
async def test_remove_config_subentry_from_device_removes_entities(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test that we remove entities tied to a device when config subentry is removed."""
config_entry_1 = MockConfigEntry(
domain="hue",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
),
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-2",
title="Mock title",
unique_id="test",
),
],
)
config_entry_1.add_to_hass(hass)
# Create device with three config subentries
device_registry.async_get_or_create(
config_entry_id=config_entry_1.entry_id,
config_subentry_id="mock-subentry-id-1",
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
device_registry.async_get_or_create(
config_entry_id=config_entry_1.entry_id,
config_subentry_id="mock-subentry-id-2",
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry_1.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
assert device_entry.config_entries == {config_entry_1.entry_id}
assert device_entry.config_subentries == {
config_entry_1.entry_id: {None, "mock-subentry-id-1", "mock-subentry-id-2"},
}
# Create one entity for each config entry
entry_1 = entity_registry.async_get_or_create(
"light",
"hue",
"1234",
config_entry=config_entry_1,
config_subentry_id="mock-subentry-id-1",
device_id=device_entry.id,
)
entry_2 = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=config_entry_1,
config_subentry_id="mock-subentry-id-2",
device_id=device_entry.id,
)
entry_3 = entity_registry.async_get_or_create(
"sensor",
"device_tracker",
"6789",
config_entry=config_entry_1,
config_subentry_id=None,
device_id=device_entry.id,
)
assert entity_registry.async_is_registered(entry_1.entity_id)
assert entity_registry.async_is_registered(entry_2.entity_id)
assert entity_registry.async_is_registered(entry_3.entity_id)
# Remove the first config subentry from the device, the entity associated with it
# should be removed
device_registry.async_update_device(
device_entry.id,
remove_config_entry_id=config_entry_1.entry_id,
remove_config_subentry_id="mock-subentry-id-1",
)
await hass.async_block_till_done()
assert device_registry.async_get(device_entry.id)
assert not entity_registry.async_is_registered(entry_1.entity_id)
assert entity_registry.async_is_registered(entry_2.entity_id)
assert entity_registry.async_is_registered(entry_3.entity_id)
# Remove the second config subentry from the device, the entity associated with it
# should be removed
device_registry.async_update_device(
device_entry.id,
remove_config_entry_id=config_entry_1.entry_id,
remove_config_subentry_id=None,
)
await hass.async_block_till_done()
assert device_registry.async_get(device_entry.id)
assert not entity_registry.async_is_registered(entry_1.entity_id)
assert entity_registry.async_is_registered(entry_2.entity_id)
assert not entity_registry.async_is_registered(entry_3.entity_id)
# Remove the third config subentry from the device, the entity associated with it
# (and the device itself) should be removed
device_registry.async_update_device(
device_entry.id,
remove_config_entry_id=config_entry_1.entry_id,
remove_config_subentry_id="mock-subentry-id-2",
)
await hass.async_block_till_done()
assert not device_registry.async_get(device_entry.id)
assert not entity_registry.async_is_registered(entry_1.entity_id)
assert not entity_registry.async_is_registered(entry_2.entity_id)
assert not entity_registry.async_is_registered(entry_3.entity_id)
async def test_remove_config_subentry_from_device_removes_entities_2(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test that we don't remove entities with no config entry when device is modified."""
config_entry_1 = MockConfigEntry(
domain="hue",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
),
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-2",
title="Mock title",
unique_id="test",
),
],
)
config_entry_1.add_to_hass(hass)
# Create device with three config subentries
device_registry.async_get_or_create(
config_entry_id=config_entry_1.entry_id,
config_subentry_id="mock-subentry-id-1",
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
device_registry.async_get_or_create(
config_entry_id=config_entry_1.entry_id,
config_subentry_id="mock-subentry-id-2",
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry_1.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
)
assert device_entry.config_entries == {config_entry_1.entry_id}
assert device_entry.config_subentries == {
config_entry_1.entry_id: {None, "mock-subentry-id-1", "mock-subentry-id-2"},
}
# Create an entity without config entry or subentry
entry_1 = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
device_id=device_entry.id,
)
assert entity_registry.async_is_registered(entry_1.entity_id)
# Remove the first config subentry from the device
device_registry.async_update_device(
device_entry.id,
remove_config_entry_id=config_entry_1.entry_id,
remove_config_subentry_id=None,
)
await hass.async_block_till_done()
assert device_registry.async_get(device_entry.id)
assert entity_registry.async_is_registered(entry_1.entity_id)
# Remove the second config subentry from the device
device_registry.async_update_device(
device_entry.id,
remove_config_entry_id=config_entry_1.entry_id,
remove_config_subentry_id="mock-subentry-id-1",
)
await hass.async_block_till_done()
assert device_registry.async_get(device_entry.id)
assert entity_registry.async_is_registered(entry_1.entity_id)
async def test_update_device_race(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
@ -1813,11 +2144,45 @@ async def test_unique_id_non_string(
)
@pytest.mark.parametrize(
("create_kwargs", "migrate_kwargs", "new_subentry_id"),
[
({}, {}, None),
({"config_subentry_id": None}, {}, None),
({}, {"new_config_subentry_id": None}, None),
({}, {"new_config_subentry_id": "mock-subentry-id-2"}, "mock-subentry-id-2"),
(
{"config_subentry_id": "mock-subentry-id-1"},
{"new_config_subentry_id": None},
None,
),
(
{"config_subentry_id": "mock-subentry-id-1"},
{"new_config_subentry_id": "mock-subentry-id-2"},
"mock-subentry-id-2",
),
],
)
def test_migrate_entity_to_new_platform(
hass: HomeAssistant, entity_registry: er.EntityRegistry
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
create_kwargs: dict,
migrate_kwargs: dict,
new_subentry_id: str | None,
) -> None:
"""Test migrate_entity_to_new_platform."""
orig_config_entry = MockConfigEntry(domain="light")
orig_config_entry = MockConfigEntry(
domain="light",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
),
],
)
orig_config_entry.add_to_hass(hass)
orig_unique_id = "5678"
orig_entry = entity_registry.async_get_or_create(
@ -1831,6 +2196,7 @@ def test_migrate_entity_to_new_platform(
original_device_class="mock-device-class",
original_icon="initial-original_icon",
original_name="initial-original_name",
**create_kwargs,
)
assert entity_registry.async_get("light.light") is orig_entry
entity_registry.async_update_entity(
@ -1839,7 +2205,18 @@ def test_migrate_entity_to_new_platform(
icon="new_icon",
)
new_config_entry = MockConfigEntry(domain="light")
new_config_entry = MockConfigEntry(
domain="light",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-2",
title="Mock title",
unique_id="test",
),
],
)
new_config_entry.add_to_hass(hass)
new_unique_id = "1234"
assert entity_registry.async_update_entity_platform(
@ -1847,6 +2224,7 @@ def test_migrate_entity_to_new_platform(
"hue2",
new_unique_id=new_unique_id,
new_config_entry_id=new_config_entry.entry_id,
**migrate_kwargs,
)
assert not entity_registry.async_get_entity_id("light", "hue", orig_unique_id)
@ -1854,6 +2232,7 @@ def test_migrate_entity_to_new_platform(
assert (new_entry := entity_registry.async_get("light.light")) is not orig_entry
assert new_entry.config_entry_id == new_config_entry.entry_id
assert new_entry.config_subentry_id == new_subentry_id
assert new_entry.unique_id == new_unique_id
assert new_entry.name == "new_name"
assert new_entry.icon == "new_icon"
@ -1886,6 +2265,97 @@ def test_migrate_entity_to_new_platform(
)
def test_migrate_entity_to_new_platform_error_handling(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
) -> None:
"""Test migrate_entity_to_new_platform."""
orig_config_entry = MockConfigEntry(
domain="light",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1",
title="Mock title",
unique_id="test",
),
],
)
orig_config_entry.add_to_hass(hass)
orig_unique_id = "5678"
orig_entry = entity_registry.async_get_or_create(
"light",
"hue",
orig_unique_id,
suggested_object_id="light",
config_entry=orig_config_entry,
config_subentry_id="mock-subentry-id-1",
disabled_by=er.RegistryEntryDisabler.USER,
entity_category=EntityCategory.CONFIG,
original_device_class="mock-device-class",
original_icon="initial-original_icon",
original_name="initial-original_name",
)
assert entity_registry.async_get("light.light") is orig_entry
new_config_entry = MockConfigEntry(
domain="light",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-2",
title="Mock title",
unique_id="test",
),
],
)
new_config_entry.add_to_hass(hass)
new_unique_id = "1234"
# Test migrating nonexisting entity
with pytest.raises(KeyError, match="'light.not_a_real_light'"):
entity_registry.async_update_entity_platform(
"light.not_a_real_light",
"hue2",
new_unique_id=new_unique_id,
new_config_entry_id=new_config_entry.entry_id,
)
# Test migrate entity without new config entry ID
with pytest.raises(
ValueError,
match="new_config_entry_id required because light.light is already linked to a config entry",
):
entity_registry.async_update_entity_platform(
"light.light",
"hue3",
)
# Test migrate entity without new config subentry ID
with pytest.raises(
ValueError,
match="Can't change config entry without changing subentry",
):
entity_registry.async_update_entity_platform(
"light.light",
"hue3",
new_config_entry_id=new_config_entry.entry_id,
)
# Test entity with a state
hass.states.async_set("light.light", "on")
with pytest.raises(
ValueError, match="Only entities that haven't been loaded can be migrated"
):
entity_registry.async_update_entity_platform(
"light.light",
"hue2",
new_unique_id=new_unique_id,
new_config_entry_id=new_config_entry.entry_id,
)
async def test_restore_entity(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
@ -1893,12 +2363,27 @@ async def test_restore_entity(
) -> None:
"""Make sure entity registry id is stable and entity_id is reused if possible."""
update_events = async_capture_events(hass, er.EVENT_ENTITY_REGISTRY_UPDATED)
config_entry = MockConfigEntry(domain="light")
config_entry = MockConfigEntry(
domain="light",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1-1",
title="Mock title",
unique_id="test",
),
],
)
config_entry.add_to_hass(hass)
entry1 = entity_registry.async_get_or_create(
"light", "hue", "1234", config_entry=config_entry
)
entry2 = entity_registry.async_get_or_create(
"light", "hue", "5678", config_entry=config_entry
"light",
"hue",
"5678",
config_entry=config_entry,
config_subentry_id="mock-subentry-id-1-1",
)
entry1 = entity_registry.async_update_entity(
@ -1922,8 +2407,11 @@ async def test_restore_entity(
# entity_id is not restored
assert attr.evolve(entry1, entity_id="light.hue_1234") == entry1_restored
assert entry2 != entry2_restored
# Config entry is not restored
assert attr.evolve(entry2, config_entry_id=None) == entry2_restored
# Config entry and subentry are not restored
assert (
attr.evolve(entry2, config_entry_id=None, config_subentry_id=None)
== entry2_restored
)
# Remove two of the entities again, then bump time
entity_registry.async_remove(entry1_restored.entity_id)
@ -2230,3 +2718,129 @@ async def test_async_remove_thread_safety(
match="Detected code that calls entity_registry.async_remove from a thread.",
):
await hass.async_add_executor_job(entity_registry.async_remove, entry.entity_id)
async def test_subentry(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
) -> None:
"""Test subentry error handling."""
entry1 = MockConfigEntry(
domain="light",
entry_id="mock-id-1",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1-1",
title="Mock title",
unique_id="test",
),
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-1-2",
title="Mock title",
unique_id="test",
),
],
)
entry1.add_to_hass(hass)
entry2 = MockConfigEntry(
domain="light",
entry_id="mock-id-2",
subentries_data=[
config_entries.ConfigSubentryData(
data={},
subentry_id="mock-subentry-id-2-1",
title="Mock title",
unique_id="test",
)
],
)
entry2.add_to_hass(hass)
with pytest.raises(
ValueError, match="Config entry mock-id-1 has no subentry bad-subentry-id"
):
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry1,
config_subentry_id="bad-subentry-id",
)
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry1,
config_subentry_id="mock-subentry-id-1-1",
)
assert entry.config_subentry_id == "mock-subentry-id-1-1"
# Try updating subentry
with pytest.raises(
ValueError, match="Config entry mock-id-1 has no subentry bad-subentry-id"
):
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry1,
config_subentry_id="bad-subentry-id",
)
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry1,
config_subentry_id="mock-subentry-id-1-2",
)
assert entry.config_subentry_id == "mock-subentry-id-1-2"
with pytest.raises(
ValueError, match="Can't change config entry without changing subentry"
):
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry2,
)
with pytest.raises(
ValueError, match="Config entry mock-id-2 has no subentry mock-subentry-id-1-2"
):
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry2,
config_subentry_id="mock-subentry-id-1-2",
)
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry2,
config_subentry_id="mock-subentry-id-2-1",
)
assert entry.config_subentry_id == "mock-subentry-id-2-1"
entry = entity_registry.async_get_or_create(
"light",
"hue",
"5678",
config_entry=entry1,
config_subentry_id=None,
)
assert entry.config_subentry_id is None
entry = entity_registry.async_update_entity(
entry.entity_id,
config_entry_id=entry2.entry_id,
config_subentry_id="mock-subentry-id-2-1",
)
assert entry.config_subentry_id == "mock-subentry-id-2-1"