Skip to content

Commit

Permalink
Migrate ESPHome unique ids to new format (home-assistant#99451)
Browse files Browse the repository at this point in the history
  • Loading branch information
bdraco authored Oct 16, 2023
1 parent 17c9d85 commit 88296c1
Show file tree
Hide file tree
Showing 5 changed files with 150 additions and 14 deletions.
10 changes: 7 additions & 3 deletions homeassistant/components/esphome/entity.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,13 @@
from collections.abc import Callable
import functools
import math
from typing import Any, Generic, TypeVar, cast
from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast

from aioesphomeapi import (
EntityCategory as EsphomeEntityCategory,
EntityInfo,
EntityState,
build_unique_id,
)
import voluptuous as vol

Expand Down Expand Up @@ -215,9 +216,12 @@ def _on_static_info_update(self, static_info: EntityInfo) -> None:
This method can be overridden in child classes to know
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._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_name = static_info.name
if entity_category := static_info.entity_category:
Expand Down
27 changes: 20 additions & 7 deletions homeassistant/components/esphome/entry_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,16 +31,19 @@
SwitchInfo,
TextSensorInfo,
UserService,
build_unique_id,
)
from aioesphomeapi.model import ButtonInfo

from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
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.storage import Store

from .bluetooth.device import ESPHomeBluetoothDevice
from .const import DOMAIN
from .dashboard import async_get_dashboard

INFO_TO_COMPONENT_TYPE: Final = {v: k for k, v in COMPONENT_TYPE_TO_INFO.items()}
Expand Down Expand Up @@ -244,24 +247,34 @@ async def _ensure_platforms_loaded(
self.loaded_platforms |= needed

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:
"""Distribute an update of static infos to all platforms."""
# First, load all platforms
needed_platforms = set()

if async_get_dashboard(hass):
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.SELECT)

ent_reg = er.async_get(hass)
registry_get_entity = ent_reg.async_get_entity_id
for info in infos:
for info_type, platform in INFO_TYPE_TO_PLATFORM.items():
if isinstance(info, info_type):
needed_platforms.add(platform)
break
platform = INFO_TYPE_TO_PLATFORM[type(info)]
needed_platforms.add(platform)
# If the unique id is in the old format, migrate it
# 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)

# Make a dict of the EntityInfo by type and send
Expand Down
9 changes: 7 additions & 2 deletions homeassistant/components/esphome/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -450,7 +450,9 @@ async def on_connect(self) -> None:

try:
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 cli.subscribe_states(entry_data.async_update_state)
await cli.subscribe_service_calls(self.async_on_service_call)
Expand Down Expand Up @@ -544,7 +546,10 @@ async def async_start(self) -> None:
self.reconnect_logic = reconnect_logic

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)

if entry_data.device_info is not None and entry_data.device_info.name:
Expand Down
110 changes: 110 additions & 0 deletions tests/components/esphome/test_entry_data.py
Original file line number Diff line number Diff line change
@@ -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"
8 changes: 6 additions & 2 deletions tests/components/esphome/test_sensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,9 @@ async def test_generic_numeric_sensor_with_entity_category_and_icon(
entity_reg = er.async_get(hass)
entry = entity_reg.async_get("sensor.test_mysensor")
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


Expand Down Expand Up @@ -152,7 +154,9 @@ async def test_generic_numeric_sensor_state_class_measurement(
entity_reg = er.async_get(hass)
entry = entity_reg.async_get("sensor.test_mysensor")
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


Expand Down

0 comments on commit 88296c1

Please sign in to comment.