From 61e4ebf15541b7729d9d80f7f044d7a2312dede0 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Thu, 11 Nov 2021 07:22:52 +0100 Subject: [PATCH] Add button entity platform (restart button) to WLED (#59086) Co-authored-by: Tom Brien --- homeassistant/components/wled/__init__.py | 10 ++- homeassistant/components/wled/button.py | 41 ++++++++++ tests/components/wled/test_button.py | 94 +++++++++++++++++++++++ 3 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 homeassistant/components/wled/button.py create mode 100644 tests/components/wled/test_button.py diff --git a/homeassistant/components/wled/__init__.py b/homeassistant/components/wled/__init__.py index e7697676014..d54cfe3cd1c 100644 --- a/homeassistant/components/wled/__init__.py +++ b/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: diff --git a/homeassistant/components/wled/button.py b/homeassistant/components/wled/button.py new file mode 100644 index 00000000000..7278495b3fa --- /dev/null +++ b/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() diff --git a/tests/components/wled/test_button.py b/tests/components/wled/test_button.py new file mode 100644 index 00000000000..9034632ae93 --- /dev/null +++ b/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