From d57b502551b62033c30e618eccdb298e80b4d80d Mon Sep 17 00:00:00 2001 From: Shay Levy Date: Tue, 7 Oct 2025 23:01:30 +0300 Subject: [PATCH] Migrate Shelly virtual button platfrom unique IDs to include roles (#153865) --- homeassistant/components/shelly/button.py | 107 +++++++++++------- homeassistant/components/shelly/entity.py | 10 +- .../shelly/snapshots/test_button.ambr | 4 +- tests/components/shelly/test_button.py | 63 ++++++++++- 4 files changed, 137 insertions(+), 47 deletions(-) diff --git a/homeassistant/components/shelly/button.py b/homeassistant/components/shelly/button.py index fbc46160f1ce..4c6d5695f335 100644 --- a/homeassistant/components/shelly/button.py +++ b/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") diff --git a/homeassistant/components/shelly/entity.py b/homeassistant/components/shelly/entity.py index 327d40df4213..44a75e64b317 100644 --- a/homeassistant/components/shelly/entity.py +++ b/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 diff --git a/tests/components/shelly/snapshots/test_button.ambr b/tests/components/shelly/snapshots/test_button.ambr index af19860f5466..7ec15e7b1db0 100644 --- a/tests/components/shelly/snapshots/test_button.ambr +++ b/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, }) # --- diff --git a/tests/components/shelly/test_button.py b/tests/components/shelly/test_button.py index f6a3df0bb486..dd1f56872e10 100644 --- a/tests/components/shelly/test_button.py +++ b/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