Migrate Shelly virtual button platfrom unique IDs to include roles (#153865)

homeassistant/components/shelly/button.py

@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING, Any, Final
 
 from aioshelly.const import BLU_TRV_IDENTIFIER, MODEL_BLU_GATEWAY_G3, RPC_GENERATIONS
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
-from aioshelly.rpc_device import RpcDevice
 
 from homeassistant.components.button import (
     DOMAIN as BUTTON_PLATFORM,
@@ -24,16 +23,24 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from .const import DOMAIN, LOGGER, SHELLY_GAS_MODELS
+from .const import DOMAIN, LOGGER, MODEL_FRANKEVER_WATER_VALVE, SHELLY_GAS_MODELS
 from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
-from .entity import get_entity_block_device_info, get_entity_rpc_device_info
+from .entity import (
+    RpcEntityDescription,
+    ShellyRpcAttributeEntity,
+    async_setup_entry_rpc,
+    get_entity_block_device_info,
+    get_entity_rpc_device_info,
+    rpc_call,
+)
 from .utils import (
     async_remove_orphaned_entities,
     format_ble_addr,
     get_blu_trv_device_info,
     get_device_entry_gen,
-    get_rpc_entity_name,
     get_rpc_key_ids,
+    get_rpc_key_instances,
+    get_rpc_role_by_key,
     get_virtual_component_ids,
 )
 
@@ -51,6 +58,11 @@ class ShellyButtonDescription[
     supported: Callable[[_ShellyCoordinatorT], bool] = lambda _: True
 
 
+@dataclass(frozen=True, kw_only=True)
+class RpcButtonDescription(RpcEntityDescription, ButtonEntityDescription):
+    """Class to describe a RPC button."""
+
+
 BUTTONS: Final[list[ShellyButtonDescription[Any]]] = [
     ShellyButtonDescription[ShellyBlockCoordinator | ShellyRpcCoordinator](
         key="reboot",
@@ -96,12 +108,24 @@ BLU_TRV_BUTTONS: Final[list[ShellyButtonDescription]] = [
     ),
 ]
 
-VIRTUAL_BUTTONS: Final[list[ShellyButtonDescription]] = [
-    ShellyButtonDescription[ShellyRpcCoordinator](
+RPC_VIRTUAL_BUTTONS = {
+    "button_generic": RpcButtonDescription(
         key="button",
-        press_action="single_push",
-    )
-]
+        role="generic",
+    ),
+    "button_open": RpcButtonDescription(
+        key="button",
+        entity_registry_enabled_default=False,
+        role="open",
+        models={MODEL_FRANKEVER_WATER_VALVE},
+    ),
+    "button_close": RpcButtonDescription(
+        key="button",
+        entity_registry_enabled_default=False,
+        role="close",
+        models={MODEL_FRANKEVER_WATER_VALVE},
+    ),
+}
 
 
 @callback
@@ -129,8 +153,10 @@ def async_migrate_unique_ids(
                 )
             }
 
+    if not isinstance(coordinator, ShellyRpcCoordinator):
+        return None
+
     if blutrv_key_ids := get_rpc_key_ids(coordinator.device.status, BLU_TRV_IDENTIFIER):
-        assert isinstance(coordinator.device, RpcDevice)
         for _id in blutrv_key_ids:
             key = f"{BLU_TRV_IDENTIFIER}:{_id}"
             ble_addr: str = coordinator.device.config[key]["addr"]
@@ -149,6 +175,26 @@ def async_migrate_unique_ids(
                     )
                 }
 
+    if virtual_button_keys := get_rpc_key_instances(
+        coordinator.device.config, "button"
+    ):
+        for key in virtual_button_keys:
+            old_unique_id = f"{coordinator.mac}-{key}"
+            if entity_entry.unique_id == old_unique_id:
+                role = get_rpc_role_by_key(coordinator.device.config, key)
+                new_unique_id = f"{coordinator.mac}-{key}-button_{role}"
+                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
 
 
@@ -172,7 +218,7 @@ async def async_setup_entry(
         hass, config_entry.entry_id, partial(async_migrate_unique_ids, coordinator)
     )
 
-    entities: list[ShellyButton | ShellyBluTrvButton | ShellyVirtualButton] = []
+    entities: list[ShellyButton | ShellyBluTrvButton] = []
 
     entities.extend(
         ShellyButton(coordinator, button)
@@ -185,12 +231,9 @@ async def async_setup_entry(
         return
 
     # add virtual buttons
-    if virtual_button_ids := get_rpc_key_ids(coordinator.device.status, "button"):
-        entities.extend(
-            ShellyVirtualButton(coordinator, button, id_)
-            for id_ in virtual_button_ids
-            for button in VIRTUAL_BUTTONS
-        )
+    async_setup_entry_rpc(
+        hass, config_entry, async_add_entities, RPC_VIRTUAL_BUTTONS, RpcVirtualButton
+    )
 
     # add BLU TRV buttons
     if blutrv_key_ids := get_rpc_key_ids(coordinator.device.status, BLU_TRV_IDENTIFIER):
@@ -332,30 +375,16 @@ class ShellyBluTrvButton(ShellyBaseButton):
         await method(self._id)
 
 
-class ShellyVirtualButton(ShellyBaseButton):
-    """Defines a Shelly virtual component button."""
+class RpcVirtualButton(ShellyRpcAttributeEntity, ButtonEntity):
+    """Defines a Shelly RPC virtual component button."""
 
-    def __init__(
-        self,
-        coordinator: ShellyRpcCoordinator,
-        description: ShellyButtonDescription,
-        _id: int,
-    ) -> None:
-        """Initialize Shelly virtual component button."""
-        super().__init__(coordinator, description)
+    entity_description: RpcButtonDescription
+    _id: int
 
-        self._attr_unique_id = f"{coordinator.mac}-{description.key}:{_id}"
-        self._attr_device_info = get_entity_rpc_device_info(coordinator)
-        self._attr_name = get_rpc_entity_name(
-            coordinator.device, f"{description.key}:{_id}"
-        )
-        self._id = _id
-
-    async def _press_method(self) -> None:
-        """Press method."""
+    @rpc_call
+    async def async_press(self) -> None:
+        """Triggers the Shelly button press service."""
         if TYPE_CHECKING:
             assert isinstance(self.coordinator, ShellyRpcCoordinator)
 
-        await self.coordinator.device.button_trigger(
-            self._id, self.entity_description.press_action
-        )
+        await self.coordinator.device.button_trigger(self._id, "single_push")

homeassistant/components/shelly/entity.py

@@ -195,9 +195,11 @@ def async_setup_rpc_attribute_entities(
             ):
                 continue
 
-            if description.sub_key not in coordinator.device.status[
-                key
-            ] and not description.supported(coordinator.device.status[key]):
+            if (
+                description.sub_key
+                and description.sub_key not in coordinator.device.status[key]
+                and not description.supported(coordinator.device.status[key])
+            ):
                 continue
 
             # Filter and remove entities that according to settings/status
@@ -309,7 +311,7 @@ class RpcEntityDescription(EntityDescription):
     # restrict the type to str.
     name: str = ""
 
-    sub_key: str
+    sub_key: str | None = None
 
     value: Callable[[Any, Any], Any] | None = None
     available: Callable[[dict], bool] | None = None

tests/components/shelly/snapshots/test_button.ambr

@@ -127,7 +127,7 @@
     'suggested_object_id': None,
     'supported_features': 0,
     'translation_key': None,
-    'unique_id': '123456789ABC-button:200',
+    'unique_id': '123456789ABC-button:200-button_generic',
     'unit_of_measurement': None,
   })
 # ---
@@ -175,7 +175,7 @@
     'suggested_object_id': None,
     'supported_features': 0,
     'translation_key': None,
-    'unique_id': '123456789ABC-button:200',
+    'unique_id': '123456789ABC-button:200-button_generic',
     'unit_of_measurement': None,
   })
 # ---

tests/components/shelly/test_button.py

@@ -9,7 +9,7 @@ import pytest
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
-from homeassistant.components.shelly.const import DOMAIN
+from homeassistant.components.shelly.const import DOMAIN, MODEL_FRANKEVER_WATER_VALVE
 from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState
 from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN, Platform
 from homeassistant.core import HomeAssistant
@@ -17,7 +17,13 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import DeviceRegistry
 from homeassistant.helpers.entity_registry import EntityRegistry
 
-from . import init_integration, patch_platforms, register_device, register_entity
+from . import (
+    MOCK_MAC,
+    init_integration,
+    patch_platforms,
+    register_device,
+    register_entity,
+)
 
 
 @pytest.fixture(autouse=True)
@@ -417,3 +423,56 @@ async def test_migrate_unique_id_blu_trv(
     assert entity_entry.unique_id == "F8447725F0DD-blutrv:200-calibrate"
 
     assert "Migrating unique_id for button.trv_name_calibrate" in caplog.text
+
+
+@pytest.mark.parametrize(
+    ("old_id", "new_id", "role"),
+    [
+        ("button", "button_generic", None),
+        ("button", "button_open", "open"),
+        ("button", "button_close", "close"),
+    ],
+)
+async def test_migrate_unique_id_virtual_components_roles(
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    entity_registry: EntityRegistry,
+    caplog: pytest.LogCaptureFixture,
+    monkeypatch: pytest.MonkeyPatch,
+    old_id: str,
+    new_id: str,
+    role: str | None,
+) -> None:
+    """Test migration of unique_id for virtual components to include role."""
+    entry = await init_integration(
+        hass, 3, model=MODEL_FRANKEVER_WATER_VALVE, skip_setup=True
+    )
+    old_unique_id = f"{MOCK_MAC}-{old_id}:200"
+    new_unique_id = f"{old_unique_id}-{new_id}"
+    config = deepcopy(mock_rpc_device.config)
+    if role:
+        config[f"{old_id}:200"] = {
+            "role": role,
+        }
+    else:
+        config[f"{old_id}:200"] = {}
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    entity = entity_registry.async_get_or_create(
+        suggested_object_id="test_name_test_button",
+        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_test_button")
+    assert entity_entry
+    assert entity_entry.unique_id == new_unique_id
+
+    assert "Migrating unique_id for button.test_name_test_button" in caplog.text