This commit is contained in:
J. Nick Koston
2025-06-24 16:43:54 +02:00
parent ea77f13ebc
commit 2446a879ca
3 changed files with 353 additions and 13 deletions

View File

@@ -13,7 +13,6 @@ from aioesphomeapi import (
EntityCategory as EsphomeEntityCategory,
EntityInfo,
EntityState,
build_unique_id,
)
import voluptuous as vol
@@ -33,7 +32,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DOMAIN
# Import config flow so that it's added to the registry
from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
from .entry_data import ESPHomeConfigEntry, RuntimeEntryData, build_device_unique_id
from .enum_mapper import EsphomeEnumMapper
_InfoT = TypeVar("_InfoT", bound=EntityInfo)
@@ -76,14 +75,21 @@ def async_static_info_updated(
if old_info.device_id == info.device_id:
continue
# Entity has switched devices, update its device assignment
unique_id = build_unique_id(device_info.mac_address, info)
entity_id = ent_reg.async_get_entity_id(platform.domain, DOMAIN, unique_id)
# entity_id should never be None here because old_info is not None,
# which means the entity was previously created and is in the registry
assert entity_id is not None
# Entity has switched devices, need to migrate unique_id
old_unique_id = build_device_unique_id(device_info.mac_address, old_info)
new_unique_id = build_device_unique_id(device_info.mac_address, info)
entity_id = ent_reg.async_get_entity_id(platform.domain, DOMAIN, old_unique_id)
# Determine the new device
if not entity_id:
continue
updates: dict[str, Any] = {}
# Update unique_id if it changed
if old_unique_id != new_unique_id:
updates["new_unique_id"] = new_unique_id
# Update device assignment
if info.device_id:
# Entity now belongs to a sub device
new_device = dev_reg.async_get_device(
@@ -96,7 +102,11 @@ def async_static_info_updated(
)
if new_device:
ent_reg.async_update_entity(entity_id, device_id=new_device.id)
updates["device_id"] = new_device.id
# Apply all updates at once
if updates:
ent_reg.async_update_entity(entity_id, **updates)
# Anything still in current_infos is now gone
if current_infos:
@@ -339,7 +349,9 @@ class EsphomeEntity(EsphomeBaseEntity, Generic[_InfoT, _StateT]):
static_info = cast(_InfoT, static_info)
assert device_info
self._static_info = static_info
self._attr_unique_id = build_unique_id(device_info.mac_address, static_info)
self._attr_unique_id = build_device_unique_id(
device_info.mac_address, static_info
)
self._attr_entity_registry_enabled_default = not static_info.disabled_by_default
# https://github.com/home-assistant/core/issues/132532
# If the name is "", we need to set it to None since otherwise

View File

@@ -95,6 +95,22 @@ INFO_TYPE_TO_PLATFORM: dict[type[EntityInfo], Platform] = {
}
def build_device_unique_id(mac: str, entity_info: EntityInfo) -> str:
"""Build unique ID for entity, appending @device_id if it belongs to a sub-device.
This wrapper around build_unique_id ensures that entities belonging to sub-devices
have their device_id appended to the unique_id to handle proper migration when
entities move between devices.
"""
base_unique_id = build_unique_id(mac, entity_info)
# If entity belongs to a sub-device, append @device_id
if entity_info.device_id:
return f"{base_unique_id}@{entity_info.device_id}"
return base_unique_id
class StoreData(TypedDict, total=False):
"""ESPHome storage data."""
@@ -223,7 +239,9 @@ class RuntimeEntryData:
ent_reg = er.async_get(hass)
for info in static_infos:
if entry := ent_reg.async_get_entity_id(
INFO_TYPE_TO_PLATFORM[type(info)], DOMAIN, build_unique_id(mac, info)
INFO_TYPE_TO_PLATFORM[type(info)],
DOMAIN,
build_device_unique_id(mac, info),
):
ent_reg.async_remove(entry)
@@ -279,7 +297,8 @@ class RuntimeEntryData:
if (
(old_unique_id := info.unique_id)
and (old_entry := registry_get_entity(platform, DOMAIN, old_unique_id))
and (new_unique_id := build_unique_id(mac, info)) != old_unique_id
and (new_unique_id := build_device_unique_id(mac, info))
!= old_unique_id
and not registry_get_entity(platform, DOMAIN, new_unique_id)
):
ent_reg.async_update_entity(old_entry, new_unique_id=new_unique_id)

View File

@@ -1162,3 +1162,312 @@ async def test_entity_id_with_empty_sub_device_name(
# When sub device has empty name, entity_id should use main device name
# Should be: binary_sensor.{main_device_name}_{object_id}
assert hass.states.get("binary_sensor.main_device_sensor") is not None
async def test_unique_id_migration_when_entity_moves_between_devices(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
device_registry: dr.DeviceRegistry,
mock_client: APIClient,
mock_esphome_device: MockESPHomeDeviceType,
) -> None:
"""Test that unique_id is migrated when entity moves between devices while entity_id stays the same."""
# Initial setup: entity on main device
device_info = {
"name": "test",
"devices": [], # No sub-devices initially
}
# Entity on main device
entity_info = [
BinarySensorInfo(
object_id="temperature",
key=1,
name="Temperature",
unique_id="unused", # This field is not used by the integration
device_id=0, # Main device
),
]
states = [
BinarySensorState(key=1, state=True, missing_state=False),
]
device = await mock_esphome_device(
mock_client=mock_client,
device_info=device_info,
entity_info=entity_info,
states=states,
)
# Check initial entity
state = hass.states.get("binary_sensor.test_temperature")
assert state is not None
# Get the entity from registry
entity_entry = entity_registry.async_get("binary_sensor.test_temperature")
assert entity_entry is not None
initial_unique_id = entity_entry.unique_id
# Initial unique_id should not have @device_id suffix since it's on main device
assert "@" not in initial_unique_id
# Add sub-device to device info
sub_devices = [
SubDeviceInfo(device_id=22222222, name="kitchen_controller", area_id=0),
]
# Get the config entry from hass
entries = hass.config_entries.async_entries(DOMAIN)
assert len(entries) == 1
entry = entries[0]
# Build device_id_to_name mapping like manager.py does
entry_data = entry.runtime_data
entry_data.device_id_to_name = {
sub_device.device_id: sub_device.name for sub_device in sub_devices
}
# Create a new DeviceInfo with sub-devices since it's frozen
# Get the current device info and convert to dict
current_device_info = mock_client.device_info.return_value
device_info_dict = asdict(current_device_info)
# Update the devices list
device_info_dict["devices"] = sub_devices
# Create new DeviceInfo with updated devices
new_device_info = DeviceInfo(**device_info_dict)
# Update mock_client to return new device info
mock_client.device_info.return_value = new_device_info
# Update entity info - same key and object_id but now on sub-device
new_entity_info = [
BinarySensorInfo(
object_id="temperature", # Same object_id
key=1, # Same key - this is what identifies the entity
name="Temperature",
unique_id="unused", # This field is not used
device_id=22222222, # Now on sub-device
),
]
# Update the entity info by changing what the mock returns
mock_client.list_entities_services = AsyncMock(return_value=(new_entity_info, []))
# Trigger a reconnect to simulate the entity info update
await device.mock_disconnect(expected_disconnect=False)
await device.mock_connect()
# Wait for entity to be updated
await hass.async_block_till_done()
# The entity_id doesn't change when moving between devices
# Only the unique_id gets updated with @device_id suffix
state = hass.states.get("binary_sensor.test_temperature")
assert state is not None
# Get updated entity from registry - entity_id should be the same
entity_entry = entity_registry.async_get("binary_sensor.test_temperature")
assert entity_entry is not None
# Unique ID should have been migrated to include @device_id
# This is done by our build_device_unique_id wrapper
expected_unique_id = f"{initial_unique_id}@22222222"
assert entity_entry.unique_id == expected_unique_id
# Entity should now be associated with the sub-device
sub_device = device_registry.async_get_device(
identifiers={(DOMAIN, f"{device.device_info.mac_address}_22222222")}
)
assert sub_device is not None
assert entity_entry.device_id == sub_device.id
async def test_unique_id_migration_sub_device_to_main_device(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
device_registry: dr.DeviceRegistry,
mock_client: APIClient,
mock_esphome_device: MockESPHomeDeviceType,
) -> None:
"""Test that unique_id is migrated when entity moves from sub-device to main device."""
# Initial setup: entity on sub-device
sub_devices = [
SubDeviceInfo(device_id=22222222, name="kitchen_controller", area_id=0),
]
device_info = {
"name": "test",
"devices": sub_devices,
}
# Entity on sub-device
entity_info = [
BinarySensorInfo(
object_id="temperature",
key=1,
name="Temperature",
unique_id="unused",
device_id=22222222, # On sub-device
),
]
states = [
BinarySensorState(key=1, state=True, missing_state=False),
]
device = await mock_esphome_device(
mock_client=mock_client,
device_info=device_info,
entity_info=entity_info,
states=states,
)
# Check initial entity
state = hass.states.get("binary_sensor.kitchen_controller_temperature")
assert state is not None
# Get the entity from registry
entity_entry = entity_registry.async_get(
"binary_sensor.kitchen_controller_temperature"
)
assert entity_entry is not None
initial_unique_id = entity_entry.unique_id
# Initial unique_id should have @device_id suffix since it's on sub-device
assert "@22222222" in initial_unique_id
# Update entity info - move to main device
new_entity_info = [
BinarySensorInfo(
object_id="temperature",
key=1,
name="Temperature",
unique_id="unused",
device_id=0, # Now on main device
),
]
# Update the entity info
mock_client.list_entities_services = AsyncMock(return_value=(new_entity_info, []))
# Trigger a reconnect
await device.mock_disconnect(expected_disconnect=False)
await device.mock_connect()
await hass.async_block_till_done()
# The entity_id should remain the same
state = hass.states.get("binary_sensor.kitchen_controller_temperature")
assert state is not None
# Get updated entity from registry
entity_entry = entity_registry.async_get(
"binary_sensor.kitchen_controller_temperature"
)
assert entity_entry is not None
# Unique ID should have been migrated to remove @device_id suffix
expected_unique_id = initial_unique_id.replace("@22222222", "")
assert entity_entry.unique_id == expected_unique_id
# Entity should now be associated with the main device
main_device = device_registry.async_get_device(
connections={(dr.CONNECTION_NETWORK_MAC, device.device_info.mac_address)}
)
assert main_device is not None
assert entity_entry.device_id == main_device.id
async def test_unique_id_migration_between_sub_devices(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
device_registry: dr.DeviceRegistry,
mock_client: APIClient,
mock_esphome_device: MockESPHomeDeviceType,
) -> None:
"""Test that unique_id is migrated when entity moves between sub-devices."""
# Initial setup: two sub-devices
sub_devices = [
SubDeviceInfo(device_id=22222222, name="kitchen_controller", area_id=0),
SubDeviceInfo(device_id=33333333, name="bedroom_controller", area_id=0),
]
device_info = {
"name": "test",
"devices": sub_devices,
}
# Entity on first sub-device
entity_info = [
BinarySensorInfo(
object_id="temperature",
key=1,
name="Temperature",
unique_id="unused",
device_id=22222222, # On kitchen_controller
),
]
states = [
BinarySensorState(key=1, state=True, missing_state=False),
]
device = await mock_esphome_device(
mock_client=mock_client,
device_info=device_info,
entity_info=entity_info,
states=states,
)
# Check initial entity
state = hass.states.get("binary_sensor.kitchen_controller_temperature")
assert state is not None
# Get the entity from registry
entity_entry = entity_registry.async_get(
"binary_sensor.kitchen_controller_temperature"
)
assert entity_entry is not None
initial_unique_id = entity_entry.unique_id
# Initial unique_id should have @22222222 suffix
assert "@22222222" in initial_unique_id
# Update entity info - move to second sub-device
new_entity_info = [
BinarySensorInfo(
object_id="temperature",
key=1,
name="Temperature",
unique_id="unused",
device_id=33333333, # Now on bedroom_controller
),
]
# Update the entity info
mock_client.list_entities_services = AsyncMock(return_value=(new_entity_info, []))
# Trigger a reconnect
await device.mock_disconnect(expected_disconnect=False)
await device.mock_connect()
await hass.async_block_till_done()
# The entity_id should remain the same
state = hass.states.get("binary_sensor.kitchen_controller_temperature")
assert state is not None
# Get updated entity from registry
entity_entry = entity_registry.async_get(
"binary_sensor.kitchen_controller_temperature"
)
assert entity_entry is not None
# Unique ID should have been migrated from @22222222 to @33333333
expected_unique_id = initial_unique_id.replace("@22222222", "@33333333")
assert entity_entry.unique_id == expected_unique_id
# Entity should now be associated with the second sub-device
bedroom_device = device_registry.async_get_device(
identifiers={(DOMAIN, f"{device.device_info.mac_address}_33333333")}
)
assert bedroom_device is not None
assert entity_entry.device_id == bedroom_device.id