Migrate ESPHome unique ids to new format (#99451)

This commit is contained in:
J. Nick Koston 2023-10-15 17:05:20 -10:00 committed by GitHub
parent 17c9d85e0e
commit 88296c1998
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 150 additions and 14 deletions

View file

@ -4,12 +4,13 @@ from __future__ import annotations
from collections.abc import Callable from collections.abc import Callable
import functools import functools
import math import math
from typing import Any, Generic, TypeVar, cast from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
from aioesphomeapi import ( from aioesphomeapi import (
EntityCategory as EsphomeEntityCategory, EntityCategory as EsphomeEntityCategory,
EntityInfo, EntityInfo,
EntityState, EntityState,
build_unique_id,
) )
import voluptuous as vol import voluptuous as vol
@ -215,9 +216,12 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
This method can be overridden in child classes to know This method can be overridden in child classes to know
when the static info changes. when the static info changes.
""" """
static_info = cast(_InfoT, static_info) device_info = self._entry_data.device_info
if TYPE_CHECKING:
static_info = cast(_InfoT, static_info)
assert device_info
self._static_info = static_info self._static_info = static_info
self._attr_unique_id = static_info.unique_id self._attr_unique_id = build_unique_id(device_info.mac_address, static_info)
self._attr_entity_registry_enabled_default = not static_info.disabled_by_default self._attr_entity_registry_enabled_default = not static_info.disabled_by_default
self._attr_name = static_info.name self._attr_name = static_info.name
if entity_category := static_info.entity_category: if entity_category := static_info.entity_category:

View file

@ -31,16 +31,19 @@ from aioesphomeapi import (
SwitchInfo, SwitchInfo,
TextSensorInfo, TextSensorInfo,
UserService, UserService,
build_unique_id,
) )
from aioesphomeapi.model import ButtonInfo from aioesphomeapi.model import ButtonInfo
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform from homeassistant.const import Platform
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.storage import Store from homeassistant.helpers.storage import Store
from .bluetooth.device import ESPHomeBluetoothDevice from .bluetooth.device import ESPHomeBluetoothDevice
from .const import DOMAIN
from .dashboard import async_get_dashboard from .dashboard import async_get_dashboard
INFO_TO_COMPONENT_TYPE: Final = {v: k for k, v in COMPONENT_TYPE_TO_INFO.items()} INFO_TO_COMPONENT_TYPE: Final = {v: k for k, v in COMPONENT_TYPE_TO_INFO.items()}
@ -244,24 +247,34 @@ class RuntimeEntryData:
self.loaded_platforms |= needed self.loaded_platforms |= needed
async def async_update_static_infos( async def async_update_static_infos(
self, hass: HomeAssistant, entry: ConfigEntry, infos: list[EntityInfo] self, hass: HomeAssistant, entry: ConfigEntry, infos: list[EntityInfo], mac: str
) -> None: ) -> None:
"""Distribute an update of static infos to all platforms.""" """Distribute an update of static infos to all platforms."""
# First, load all platforms # First, load all platforms
needed_platforms = set() needed_platforms = set()
if async_get_dashboard(hass): if async_get_dashboard(hass):
needed_platforms.add(Platform.UPDATE) needed_platforms.add(Platform.UPDATE)
if self.device_info is not None and self.device_info.voice_assistant_version: if self.device_info and self.device_info.voice_assistant_version:
needed_platforms.add(Platform.BINARY_SENSOR) needed_platforms.add(Platform.BINARY_SENSOR)
needed_platforms.add(Platform.SELECT) needed_platforms.add(Platform.SELECT)
ent_reg = er.async_get(hass)
registry_get_entity = ent_reg.async_get_entity_id
for info in infos: for info in infos:
for info_type, platform in INFO_TYPE_TO_PLATFORM.items(): platform = INFO_TYPE_TO_PLATFORM[type(info)]
if isinstance(info, info_type): needed_platforms.add(platform)
needed_platforms.add(platform) # If the unique id is in the old format, migrate it
break # except if they downgraded and upgraded, there might be a duplicate
# so we want to keep the one that was already there.
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 not registry_get_entity(platform, DOMAIN, new_unique_id)
):
ent_reg.async_update_entity(old_entry, new_unique_id=new_unique_id)
await self._ensure_platforms_loaded(hass, entry, needed_platforms) await self._ensure_platforms_loaded(hass, entry, needed_platforms)
# Make a dict of the EntityInfo by type and send # Make a dict of the EntityInfo by type and send

View file

@ -450,7 +450,9 @@ class ESPHomeManager:
try: try:
entity_infos, services = await cli.list_entities_services() entity_infos, services = await cli.list_entities_services()
await entry_data.async_update_static_infos(hass, entry, entity_infos) await entry_data.async_update_static_infos(
hass, entry, entity_infos, device_info.mac_address
)
await _setup_services(hass, entry_data, services) await _setup_services(hass, entry_data, services)
await cli.subscribe_states(entry_data.async_update_state) await cli.subscribe_states(entry_data.async_update_state)
await cli.subscribe_service_calls(self.async_on_service_call) await cli.subscribe_service_calls(self.async_on_service_call)
@ -544,7 +546,10 @@ class ESPHomeManager:
self.reconnect_logic = reconnect_logic self.reconnect_logic = reconnect_logic
infos, services = await entry_data.async_load_from_store() infos, services = await entry_data.async_load_from_store()
await entry_data.async_update_static_infos(hass, entry, infos) if entry.unique_id:
await entry_data.async_update_static_infos(
hass, entry, infos, entry.unique_id.upper()
)
await _setup_services(hass, entry_data, services) await _setup_services(hass, entry_data, services)
if entry_data.device_info is not None and entry_data.device_info.name: if entry_data.device_info is not None and entry_data.device_info.name:

View file

@ -0,0 +1,110 @@
"""Test ESPHome entry data."""
from aioesphomeapi import (
APIClient,
EntityCategory as ESPHomeEntityCategory,
SensorInfo,
SensorState,
)
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
async def test_migrate_entity_unique_id(
hass: HomeAssistant,
mock_client: APIClient,
mock_generic_device_entry,
) -> None:
"""Test a generic sensor entity unique id migration."""
ent_reg = er.async_get(hass)
ent_reg.async_get_or_create(
"sensor",
"esphome",
"my_sensor",
suggested_object_id="old_sensor",
disabled_by=None,
)
entity_info = [
SensorInfo(
object_id="mysensor",
key=1,
name="my sensor",
unique_id="my_sensor",
entity_category=ESPHomeEntityCategory.DIAGNOSTIC,
icon="mdi:leaf",
)
]
states = [SensorState(key=1, state=50)]
user_service = []
await mock_generic_device_entry(
mock_client=mock_client,
entity_info=entity_info,
user_service=user_service,
states=states,
)
state = hass.states.get("sensor.old_sensor")
assert state is not None
assert state.state == "50"
entity_reg = er.async_get(hass)
entry = entity_reg.async_get("sensor.old_sensor")
assert entry is not None
assert entity_reg.async_get_entity_id("sensor", "esphome", "my_sensor") is None
# Note that ESPHome includes the EntityInfo type in the unique id
# as this is not a 1:1 mapping to the entity platform (ie. text_sensor)
assert entry.unique_id == "11:22:33:44:55:aa-sensor-mysensor"
async def test_migrate_entity_unique_id_downgrade_upgrade(
hass: HomeAssistant,
mock_client: APIClient,
mock_generic_device_entry,
) -> None:
"""Test unique id migration prefers the original entity on downgrade upgrade."""
ent_reg = er.async_get(hass)
ent_reg.async_get_or_create(
"sensor",
"esphome",
"my_sensor",
suggested_object_id="old_sensor",
disabled_by=None,
)
ent_reg.async_get_or_create(
"sensor",
"esphome",
"11:22:33:44:55:aa-sensor-mysensor",
suggested_object_id="new_sensor",
disabled_by=None,
)
entity_info = [
SensorInfo(
object_id="mysensor",
key=1,
name="my sensor",
unique_id="my_sensor",
entity_category=ESPHomeEntityCategory.DIAGNOSTIC,
icon="mdi:leaf",
)
]
states = [SensorState(key=1, state=50)]
user_service = []
await mock_generic_device_entry(
mock_client=mock_client,
entity_info=entity_info,
user_service=user_service,
states=states,
)
state = hass.states.get("sensor.new_sensor")
assert state is not None
assert state.state == "50"
entity_reg = er.async_get(hass)
entry = entity_reg.async_get("sensor.new_sensor")
assert entry is not None
# Confirm we did not touch the entity that was created
# on downgrade so when they upgrade again they can delete the
# entity that was only created on downgrade and they keep
# the original one.
assert entity_reg.async_get_entity_id("sensor", "esphome", "my_sensor") is not None
# Note that ESPHome includes the EntityInfo type in the unique id
# as this is not a 1:1 mapping to the entity platform (ie. text_sensor)
assert entry.unique_id == "11:22:33:44:55:aa-sensor-mysensor"

View file

@ -116,7 +116,9 @@ async def test_generic_numeric_sensor_with_entity_category_and_icon(
entity_reg = er.async_get(hass) entity_reg = er.async_get(hass)
entry = entity_reg.async_get("sensor.test_mysensor") entry = entity_reg.async_get("sensor.test_mysensor")
assert entry is not None assert entry is not None
assert entry.unique_id == "my_sensor" # Note that ESPHome includes the EntityInfo type in the unique id
# as this is not a 1:1 mapping to the entity platform (ie. text_sensor)
assert entry.unique_id == "11:22:33:44:55:aa-sensor-mysensor"
assert entry.entity_category is EntityCategory.DIAGNOSTIC assert entry.entity_category is EntityCategory.DIAGNOSTIC
@ -152,7 +154,9 @@ async def test_generic_numeric_sensor_state_class_measurement(
entity_reg = er.async_get(hass) entity_reg = er.async_get(hass)
entry = entity_reg.async_get("sensor.test_mysensor") entry = entity_reg.async_get("sensor.test_mysensor")
assert entry is not None assert entry is not None
assert entry.unique_id == "my_sensor" # Note that ESPHome includes the EntityInfo type in the unique id
# as this is not a 1:1 mapping to the entity platform (ie. text_sensor)
assert entry.unique_id == "11:22:33:44:55:aa-sensor-mysensor"
assert entry.entity_category is None assert entry.entity_category is None