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

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

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,