Add reboot button to Magic Home/flux_led (#62323)

homeassistant/components/flux_led/__init__.py

@@ -40,8 +40,13 @@ from .discovery import (
 _LOGGER = logging.getLogger(__name__)
 
 PLATFORMS_BY_TYPE: Final = {
-    DeviceType.Bulb: [Platform.LIGHT, Platform.NUMBER, Platform.SWITCH],
+    DeviceType.Bulb: [
-    DeviceType.Switch: [Platform.SWITCH],
+        Platform.BUTTON,
+        Platform.LIGHT,
+        Platform.NUMBER,
+        Platform.SWITCH,
+    ],
+    DeviceType.Switch: [Platform.BUTTON, Platform.SWITCH],
 }
 DISCOVERY_INTERVAL: Final = timedelta(minutes=15)
 REQUEST_REFRESH_DELAY: Final = 1.5

homeassistant/components/flux_led/button.py

@@ -0,0 +1,47 @@
+"""Support for Magic home button."""
+from __future__ import annotations
+
+from flux_led.aio import AIOWifiLedBulb
+
+from homeassistant import config_entries
+from homeassistant.components.button import ButtonEntity
+from homeassistant.const import CONF_NAME
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from . import FluxLedUpdateCoordinator
+from .const import DOMAIN
+from .entity import FluxBaseEntity
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: config_entries.ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up Magic Home button based on a config entry."""
+    coordinator: FluxLedUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities([FluxRestartButton(coordinator.device, entry)])
+
+
+class FluxRestartButton(FluxBaseEntity, ButtonEntity):
+    """Representation of a Flux restart button."""
+
+    _attr_should_poll = False
+    _attr_entity_category = EntityCategory.CONFIG
+
+    def __init__(
+        self,
+        device: AIOWifiLedBulb,
+        entry: config_entries.ConfigEntry,
+    ) -> None:
+        """Initialize the reboot button."""
+        super().__init__(device, entry)
+        self._attr_name = f"{entry.data[CONF_NAME]} Restart"
+        if entry.unique_id:
+            self._attr_unique_id = f"{entry.unique_id}_restart"
+
+    async def async_press(self) -> None:
+        """Send out a restart command."""
+        await self._device.async_reboot()

tests/components/flux_led/__init__.py

@@ -124,6 +124,7 @@ def _mocked_switch() -> AIOWifiLedBulb:
 
     switch.device_type = DeviceType.Switch
     switch.requires_turn_on = True
+    switch.async_reboot = AsyncMock()
     switch.async_setup = AsyncMock(side_effect=_save_setup_callback)
     switch.async_stop = AsyncMock()
     switch.async_update = AsyncMock()

tests/components/flux_led/test_button.py

@@ -0,0 +1,41 @@
+"""Tests for button platform."""
+from homeassistant.components import flux_led
+from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN
+from homeassistant.components.flux_led.const import DOMAIN
+from homeassistant.const import ATTR_ENTITY_ID, CONF_HOST, CONF_NAME
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from . import (
+    DEFAULT_ENTRY_TITLE,
+    IP_ADDRESS,
+    MAC_ADDRESS,
+    _mocked_switch,
+    _patch_discovery,
+    _patch_wifibulb,
+)
+
+from tests.common import MockConfigEntry
+
+
+async def test_switch_reboot(hass: HomeAssistant) -> None:
+    """Test a smart plug can be rebooted."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_HOST: IP_ADDRESS, CONF_NAME: DEFAULT_ENTRY_TITLE},
+        unique_id=MAC_ADDRESS,
+    )
+    config_entry.add_to_hass(hass)
+    switch = _mocked_switch()
+    with _patch_discovery(), _patch_wifibulb(device=switch):
+        await async_setup_component(hass, flux_led.DOMAIN, {flux_led.DOMAIN: {}})
+        await hass.async_block_till_done()
+
+    entity_id = "button.bulb_rgbcw_ddeeff_restart"
+
+    assert hass.states.get(entity_id)
+
+    await hass.services.async_call(
+        BUTTON_DOMAIN, "press", {ATTR_ENTITY_ID: entity_id}, blocking=True
+    )
+    switch.async_reboot.assert_called_once()