Do not create restart button for sleeping gen2+ Shelly devices (home-assistant#158047)

Author: bieniu

homeassistant/components/shelly/button.py

@@ -44,6 +44,7 @@
 )
 from .utils import (
     async_remove_orphaned_entities,
+    async_remove_shelly_entity,
     format_ble_addr,
     get_blu_trv_device_info,
     get_device_entry_gen,
@@ -80,6 +81,7 @@ class RpcButtonDescription(RpcEntityDescription, ButtonEntityDescription):
         device_class=ButtonDeviceClass.RESTART,
         entity_category=EntityCategory.CONFIG,
         press_action="trigger_reboot",
+        supported=lambda coordinator: coordinator.sleep_period == 0,
     ),
     ShellyButtonDescription[ShellyBlockCoordinator](
         key="self_test",
@@ -197,7 +199,8 @@ async def async_setup_entry(
     """Set up button entities."""
     entry_data = config_entry.runtime_data
     coordinator: ShellyRpcCoordinator | ShellyBlockCoordinator | None
-    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
+    device_gen = get_device_entry_gen(config_entry)
+    if device_gen in RPC_GENERATIONS:
         coordinator = entry_data.rpc
     else:
         coordinator = entry_data.block
@@ -210,6 +213,12 @@ async def async_setup_entry(
             hass, config_entry.entry_id, partial(async_migrate_unique_ids, coordinator)
         )
 
+    # Remove the 'restart' button for sleeping devices as it was mistakenly
+    # added in https://github.com/home-assistant/core/pull/154673
+    entry_sleep_period = config_entry.data[CONF_SLEEP_PERIOD]
+    if device_gen in RPC_GENERATIONS and entry_sleep_period:
+        async_remove_shelly_entity(hass, BUTTON_PLATFORM, f"{coordinator.mac}-reboot")
+
     entities: list[ShellyButton] = []
 
     entities.extend(
@@ -224,7 +233,7 @@ async def async_setup_entry(
         return
 
     # add RPC buttons
-    if config_entry.data[CONF_SLEEP_PERIOD]:
+    if entry_sleep_period:
         async_setup_entry_rpc(
             hass,
             config_entry,

tests/components/shelly/test_button.py

@@ -554,3 +554,30 @@ async def test_wall_display_screen_buttons(
         blocking=True,
     )
     mock_rpc_device.wall_display_set_screen.assert_called_once_with(value=value)
+
+
+async def test_rpc_remove_restart_button_for_sleeping_devices(
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+    device_registry: DeviceRegistry,
+    entity_registry: EntityRegistry,
+) -> None:
+    """Test RPC remove restart button for sleeping devices."""
+    config_entry = await init_integration(hass, 2, sleep_period=1000, skip_setup=True)
+    device_entry = register_device(device_registry, config_entry)
+    entity_id = register_entity(
+        hass,
+        BUTTON_DOMAIN,
+        "test_name_restart",
+        "reboot",
+        config_entry,
+        device_id=device_entry.id,
+    )
+
+    assert entity_registry.async_get(entity_id) is not None
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entity_registry.async_get(entity_id) is None