Add button entity platform (restart button) to WLED (#59086)

Co-authored-by: Tom Brien <TomBrien@users.noreply.github.com>

homeassistant/components/wled/__init__.py

@@ -1,6 +1,7 @@
 """Support for WLED."""
 from __future__ import annotations
 
+from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 from homeassistant.components.number import DOMAIN as NUMBER_DOMAIN
 from homeassistant.components.select import DOMAIN as SELECT_DOMAIN
@@ -12,7 +13,14 @@ from homeassistant.core import HomeAssistant
 from .const import DOMAIN
 from .coordinator import WLEDDataUpdateCoordinator
 
-PLATFORMS = (LIGHT_DOMAIN, SELECT_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN, NUMBER_DOMAIN)
+PLATFORMS = (
+    BUTTON_DOMAIN,
+    LIGHT_DOMAIN,
+    SELECT_DOMAIN,
+    SENSOR_DOMAIN,
+    SWITCH_DOMAIN,
+    NUMBER_DOMAIN,
+)
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:

homeassistant/components/wled/button.py

@@ -0,0 +1,41 @@
+"""Support for WLED button."""
+from __future__ import annotations
+
+from homeassistant.components.button import ButtonEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import ENTITY_CATEGORY_CONFIG
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import DOMAIN
+from .coordinator import WLEDDataUpdateCoordinator
+from .helpers import wled_exception_handler
+from .models import WLEDEntity
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up WLED button based on a config entry."""
+    coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities([WLEDRestartButton(coordinator)])
+
+
+class WLEDRestartButton(WLEDEntity, ButtonEntity):
+    """Defines a WLED restart switch."""
+
+    _attr_icon = "mdi:restart"
+    _attr_entity_category = ENTITY_CATEGORY_CONFIG
+
+    def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
+        """Initialize the button entity."""
+        super().__init__(coordinator=coordinator)
+        self._attr_name = f"{coordinator.data.info.name} Restart"
+        self._attr_unique_id = f"{coordinator.data.info.mac_address}_restart"
+
+    @wled_exception_handler
+    async def async_press(self) -> None:
+        """Send out a restart command."""
+        await self.coordinator.wled.reset()

tests/components/wled/test_button.py

@@ -0,0 +1,94 @@
+"""Tests for the WLED button platform."""
+from unittest.mock import MagicMock
+
+from freezegun import freeze_time
+import pytest
+from wled import WLEDConnectionError, WLEDError
+
+from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_ICON,
+    ENTITY_CATEGORY_CONFIG,
+    STATE_UNAVAILABLE,
+    STATE_UNKNOWN,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+
+from tests.common import MockConfigEntry
+
+
+async def test_button(
+    hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
+) -> None:
+    """Test the creation and values of the WLED button."""
+    entity_registry = er.async_get(hass)
+
+    state = hass.states.get("button.wled_rgb_light_restart")
+    assert state
+    assert state.attributes.get(ATTR_ICON) == "mdi:restart"
+    assert state.state == STATE_UNKNOWN
+
+    entry = entity_registry.async_get("button.wled_rgb_light_restart")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_restart"
+    assert entry.entity_category == ENTITY_CATEGORY_CONFIG
+
+    # Restart
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.wled_rgb_light_restart"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+    assert mock_wled.reset.call_count == 1
+    mock_wled.reset.assert_called_with()
+
+
+@freeze_time("2021-11-04 17:37:00", tz_offset=-1)
+async def test_button_error(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test error handling of the WLED buttons."""
+    mock_wled.reset.side_effect = WLEDError
+
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.wled_rgb_light_restart"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("button.wled_rgb_light_restart")
+    assert state
+    assert state.state == "2021-11-04T16:37:00+00:00"
+    assert "Invalid response from API" in caplog.text
+
+
+async def test_button_connection_error(
+    hass: HomeAssistant,
+    init_integration: MockConfigEntry,
+    mock_wled: MagicMock,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test error handling of the WLED buttons."""
+    mock_wled.reset.side_effect = WLEDConnectionError
+
+    await hass.services.async_call(
+        BUTTON_DOMAIN,
+        SERVICE_PRESS,
+        {ATTR_ENTITY_ID: "button.wled_rgb_light_restart"},
+        blocking=True,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("button.wled_rgb_light_restart")
+    assert state
+    assert state.state == STATE_UNAVAILABLE
+    assert "Error communicating with API" in caplog.text