From 917b467b8591dec413fe27a04fe367d0cc2e95fd Mon Sep 17 00:00:00 2001 From: TimL Date: Thu, 22 May 2025 22:50:22 +1000 Subject: [PATCH] Add SMLIGHT button entities for second radio (#141463) * Add button entities for second radio * Update tests for second router reconnect button --- homeassistant/components/smlight/button.py | 48 ++++++++++++++-------- tests/components/smlight/test_button.py | 37 ++++++++++++++--- 2 files changed, 62 insertions(+), 23 deletions(-) diff --git a/homeassistant/components/smlight/button.py b/homeassistant/components/smlight/button.py index f834392ea13..67d9997a105 100644 --- a/homeassistant/components/smlight/button.py +++ b/homeassistant/components/smlight/button.py @@ -32,7 +32,7 @@ _LOGGER = logging.getLogger(__name__) class SmButtonDescription(ButtonEntityDescription): """Class to describe a Button entity.""" - press_fn: Callable[[CmdWrapper], Awaitable[None]] + press_fn: Callable[[CmdWrapper, int], Awaitable[None]] BUTTONS: list[SmButtonDescription] = [ @@ -40,19 +40,19 @@ BUTTONS: list[SmButtonDescription] = [ key="core_restart", translation_key="core_restart", device_class=ButtonDeviceClass.RESTART, - press_fn=lambda cmd: cmd.reboot(), + press_fn=lambda cmd, idx: cmd.reboot(), ), SmButtonDescription( key="zigbee_restart", translation_key="zigbee_restart", device_class=ButtonDeviceClass.RESTART, - press_fn=lambda cmd: cmd.zb_restart(), + press_fn=lambda cmd, idx: cmd.zb_restart(), ), SmButtonDescription( key="zigbee_flash_mode", translation_key="zigbee_flash_mode", entity_registry_enabled_default=False, - press_fn=lambda cmd: cmd.zb_bootloader(), + press_fn=lambda cmd, idx: cmd.zb_bootloader(), ), ] @@ -60,7 +60,7 @@ ROUTER = SmButtonDescription( key="reconnect_zigbee_router", translation_key="reconnect_zigbee_router", entity_registry_enabled_default=False, - press_fn=lambda cmd: cmd.zb_router(), + press_fn=lambda cmd, idx: cmd.zb_router(idx=idx), ) @@ -71,23 +71,32 @@ async def async_setup_entry( ) -> None: """Set up SMLIGHT buttons based on a config entry.""" coordinator = entry.runtime_data.data + radios = coordinator.data.info.radios async_add_entities(SmButton(coordinator, button) for button in BUTTONS) - entity_created = False + entity_created = [False, False] @callback def _check_router(startup: bool = False) -> None: - nonlocal entity_created + def router_entity(router: SmButtonDescription, idx: int) -> None: + nonlocal entity_created + zb_type = coordinator.data.info.radios[idx].zb_type - if coordinator.data.info.zb_type == 1 and not entity_created: - async_add_entities([SmButton(coordinator, ROUTER)]) - entity_created = True - elif coordinator.data.info.zb_type != 1 and (startup or entity_created): - entity_registry = er.async_get(hass) - if entity_id := entity_registry.async_get_entity_id( - BUTTON_DOMAIN, DOMAIN, f"{coordinator.unique_id}-{ROUTER.key}" - ): - entity_registry.async_remove(entity_id) + if zb_type == 1 and not entity_created[idx]: + async_add_entities([SmButton(coordinator, router, idx)]) + entity_created[idx] = True + elif zb_type != 1 and (startup or entity_created[idx]): + entity_registry = er.async_get(hass) + button = f"_{idx}" if idx else "" + if entity_id := entity_registry.async_get_entity_id( + BUTTON_DOMAIN, + DOMAIN, + f"{coordinator.unique_id}-{router.key}{button}", + ): + entity_registry.async_remove(entity_id) + + for idx, _ in enumerate(radios): + router_entity(ROUTER, idx) coordinator.async_add_listener(_check_router) _check_router(startup=True) @@ -104,13 +113,16 @@ class SmButton(SmEntity, ButtonEntity): self, coordinator: SmDataUpdateCoordinator, description: SmButtonDescription, + idx: int = 0, ) -> None: """Initialize SLZB-06 button entity.""" super().__init__(coordinator) self.entity_description = description - self._attr_unique_id = f"{coordinator.unique_id}-{description.key}" + self.idx = idx + button = f"_{idx}" if idx else "" + self._attr_unique_id = f"{coordinator.unique_id}-{description.key}{button}" async def async_press(self) -> None: """Trigger button press.""" - await self.entity_description.press_fn(self.coordinator.client.cmds) + await self.entity_description.press_fn(self.coordinator.client.cmds, self.idx) diff --git a/tests/components/smlight/test_button.py b/tests/components/smlight/test_button.py index 51e9414c00e..f9ea010fe7c 100644 --- a/tests/components/smlight/test_button.py +++ b/tests/components/smlight/test_button.py @@ -3,18 +3,22 @@ from unittest.mock import MagicMock from freezegun.api import FrozenDateTimeFactory -from pysmlight import Info +from pysmlight import Info, Radio import pytest from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS -from homeassistant.components.smlight.const import SCAN_INTERVAL +from homeassistant.components.smlight.const import DOMAIN, SCAN_INTERVAL from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN, Platform from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er from .conftest import setup_integration -from tests.common import MockConfigEntry, async_fire_time_changed +from tests.common import ( + MockConfigEntry, + async_fire_time_changed, + load_json_object_fixture, +) @pytest.fixture @@ -23,7 +27,7 @@ def platforms() -> Platform | list[Platform]: return [Platform.BUTTON] -MOCK_ROUTER = Info(MAC="AA:BB:CC:DD:EE:FF", zb_type=1) +MOCK_ROUTER = Info(MAC="AA:BB:CC:DD:EE:FF", radios=[Radio(zb_type=1)]) @pytest.mark.parametrize( @@ -67,7 +71,7 @@ async def test_buttons( ) assert len(mock_method.mock_calls) == 1 - mock_method.assert_called_with() + mock_method.assert_called() @pytest.mark.parametrize("entity_id", ["zigbee_flash_mode", "reconnect_zigbee_router"]) @@ -90,6 +94,29 @@ async def test_disabled_by_default_buttons( assert entry.disabled_by is er.RegistryEntryDisabler.INTEGRATION +@pytest.mark.usefixtures("entity_registry_enabled_by_default") +async def test_zigbee2_router_button( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + mock_config_entry: MockConfigEntry, + mock_smlight_client: MagicMock, +) -> None: + """Test creation of second radio router button (if available).""" + mock_smlight_client.get_info.side_effect = None + mock_smlight_client.get_info.return_value = Info.from_dict( + load_json_object_fixture("info-MR1.json", DOMAIN) + ) + await setup_integration(hass, mock_config_entry) + + state = hass.states.get("button.mock_title_reconnect_zigbee_router") + assert state is not None + assert state.state == STATE_UNKNOWN + + entry = entity_registry.async_get("button.mock_title_reconnect_zigbee_router") + assert entry is not None + assert entry.unique_id == "aa:bb:cc:dd:ee:ff-reconnect_zigbee_router_1" + + async def test_remove_router_reconnect( hass: HomeAssistant, entity_registry: er.EntityRegistry,