From 8d6a711cace48ec75b3e09022c09bf25160d6913 Mon Sep 17 00:00:00 2001 From: Maciej Bieniek Date: Tue, 27 Jun 2023 17:10:12 +0000 Subject: [PATCH] Make `unique_id` of the Shelly button entity immutable (#95160) --- homeassistant/components/shelly/button.py | 53 +++++++++++++++++++++-- tests/components/shelly/test_button.py | 49 +++++++++++++++++++++ 2 files changed, 98 insertions(+), 4 deletions(-) diff --git a/homeassistant/components/shelly/button.py b/homeassistant/components/shelly/button.py index 15a05573656..1f684ce137c 100644 --- a/homeassistant/components/shelly/button.py +++ b/homeassistant/components/shelly/button.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Callable, Coroutine from dataclasses import dataclass -from typing import Any, Final, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar from homeassistant.components.button import ( ButtonDeviceClass, @@ -12,14 +12,15 @@ from homeassistant.components.button import ( ) from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory -from homeassistant.core import HomeAssistant +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers import entity_registry as er from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import slugify -from .const import SHELLY_GAS_MODELS +from .const import LOGGER, SHELLY_GAS_MODELS from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .utils import get_device_entry_gen @@ -79,12 +80,52 @@ BUTTONS: Final[list[ShellyButtonDescription[Any]]] = [ ] +@callback +def async_migrate_unique_ids( + entity_entry: er.RegistryEntry, + coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator, +) -> dict[str, Any] | None: + """Migrate button unique IDs.""" + if not entity_entry.entity_id.startswith("button"): + return None + + device_name = slugify(coordinator.device.name) + + for key in ("reboot", "self_test", "mute", "unmute"): + if entity_entry.unique_id.startswith(device_name): + old_unique_id = entity_entry.unique_id + new_unique_id = f"{coordinator.mac}_{key}" + LOGGER.debug( + "Migrating unique_id for %s entity from [%s] to [%s]", + entity_entry.entity_id, + old_unique_id, + new_unique_id, + ) + return { + "new_unique_id": entity_entry.unique_id.replace( + old_unique_id, new_unique_id + ) + } + + return None + + async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set buttons for device.""" + + @callback + def _async_migrate_unique_ids( + entity_entry: er.RegistryEntry, + ) -> dict[str, Any] | None: + """Migrate button unique IDs.""" + if TYPE_CHECKING: + assert coordinator is not None + return async_migrate_unique_ids(entity_entry, coordinator) + coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None = None if get_device_entry_gen(config_entry) == 2: coordinator = get_entry_data(hass)[config_entry.entry_id].rpc @@ -92,6 +133,10 @@ async def async_setup_entry( coordinator = get_entry_data(hass)[config_entry.entry_id].block if coordinator is not None: + await er.async_migrate_entries( + hass, config_entry.entry_id, _async_migrate_unique_ids + ) + entities: list[ShellyButton] = [] for button in BUTTONS: @@ -123,7 +168,7 @@ class ShellyButton( self.entity_description = description self._attr_name = f"{coordinator.device.name} {description.name}" - self._attr_unique_id = slugify(self._attr_name) + self._attr_unique_id = f"{coordinator.mac}_{description.key}" self._attr_device_info = DeviceInfo( connections={(CONNECTION_NETWORK_MAC, coordinator.mac)} ) diff --git a/tests/components/shelly/test_button.py b/tests/components/shelly/test_button.py index fdb2cd60450..42fa83b32a1 100644 --- a/tests/components/shelly/test_button.py +++ b/tests/components/shelly/test_button.py @@ -1,9 +1,13 @@ """Tests for Shelly button platform.""" from __future__ import annotations +import pytest + from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS +from homeassistant.components.shelly.const import DOMAIN from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er from . import init_integration @@ -38,3 +42,48 @@ async def test_rpc_button(hass: HomeAssistant, mock_rpc_device) -> None: blocking=True, ) assert mock_rpc_device.trigger_reboot.call_count == 1 + + +@pytest.mark.parametrize( + ("gen", "old_unique_id", "new_unique_id", "migration"), + [ + (2, "test_name_reboot", "123456789ABC_reboot", True), + (1, "test_name_reboot", "123456789ABC_reboot", True), + (2, "123456789ABC_reboot", "123456789ABC_reboot", False), + ], +) +async def test_migrate_unique_id( + hass: HomeAssistant, + mock_block_device, + mock_rpc_device, + caplog: pytest.LogCaptureFixture, + gen: int, + old_unique_id: str, + new_unique_id: str, + migration: bool, +) -> None: + """Test migration of unique_id.""" + entry = await init_integration(hass, gen, skip_setup=True) + + entity_registry = er.async_get(hass) + entity: er.RegistryEntry = entity_registry.async_get_or_create( + suggested_object_id="test_name_reboot", + disabled_by=None, + domain=BUTTON_DOMAIN, + platform=DOMAIN, + unique_id=old_unique_id, + config_entry=entry, + ) + assert entity.unique_id == old_unique_id + + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + + entity_entry = entity_registry.async_get("button.test_name_reboot") + assert entity_entry + assert entity_entry.unique_id == new_unique_id + + assert ( + bool("Migrating unique_id for button.test_name_reboot" in caplog.text) + == migration + )