Make unique_id of the Shelly button entity immutable (#95160)

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)}
         )

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
+    )