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

This commit is contained in:
Maciej Bieniek 2023-06-27 17:10:12 +00:00 committed by GitHub
parent 1dc9c77a3e
commit 8d6a711cac
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 98 additions and 4 deletions

View File

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

View File

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