Add button device classes to WLED (#60613)

This commit is contained in:
Franck Nijhof 2021-11-30 12:12:08 +01:00 committed by GitHub
parent 9b9801516b
commit 1b8eba0afd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 42 additions and 38 deletions

View File

@ -1,7 +1,7 @@
"""Support for WLED button.""" """Support for WLED button."""
from __future__ import annotations from __future__ import annotations
from homeassistant.components.button import ButtonEntity from homeassistant.components.button import ButtonDeviceClass, ButtonEntity
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ENTITY_CATEGORY_CONFIG from homeassistant.const import ENTITY_CATEGORY_CONFIG
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
@ -23,7 +23,7 @@ async def async_setup_entry(
async_add_entities( async_add_entities(
[ [
WLEDRestartButton(coordinator), WLEDRestartButton(coordinator),
WLEDUpgradeButton(coordinator), WLEDUpdateButton(coordinator),
] ]
) )
@ -31,7 +31,7 @@ async def async_setup_entry(
class WLEDRestartButton(WLEDEntity, ButtonEntity): class WLEDRestartButton(WLEDEntity, ButtonEntity):
"""Defines a WLED restart button.""" """Defines a WLED restart button."""
_attr_icon = "mdi:restart" _attr_device_class = ButtonDeviceClass.RESTART
_attr_entity_category = ENTITY_CATEGORY_CONFIG _attr_entity_category = ENTITY_CATEGORY_CONFIG
def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None: def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
@ -46,21 +46,21 @@ class WLEDRestartButton(WLEDEntity, ButtonEntity):
await self.coordinator.wled.reset() await self.coordinator.wled.reset()
class WLEDUpgradeButton(WLEDEntity, ButtonEntity): class WLEDUpdateButton(WLEDEntity, ButtonEntity):
"""Defines a WLED upgrade button.""" """Defines a WLED update button."""
_attr_icon = "mdi:cellphone-arrow-down" _attr_device_class = ButtonDeviceClass.UPDATE
_attr_entity_category = ENTITY_CATEGORY_CONFIG _attr_entity_category = ENTITY_CATEGORY_CONFIG
def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None: def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
"""Initialize the button entity.""" """Initialize the button entity."""
super().__init__(coordinator=coordinator) super().__init__(coordinator=coordinator)
self._attr_name = f"{coordinator.data.info.name} Upgrade" self._attr_name = f"{coordinator.data.info.name} Update"
self._attr_unique_id = f"{coordinator.data.info.mac_address}_upgrade" self._attr_unique_id = f"{coordinator.data.info.mac_address}_update"
@property @property
def available(self) -> bool: def available(self) -> bool:
"""Return if the entity and an upgrade is available.""" """Return if the entity and an update is available."""
current = self.coordinator.data.info.version current = self.coordinator.data.info.version
beta = self.coordinator.data.info.version_latest_beta beta = self.coordinator.data.info.version_latest_beta
stable = self.coordinator.data.info.version_latest_stable stable = self.coordinator.data.info.version_latest_stable
@ -82,13 +82,13 @@ class WLEDUpgradeButton(WLEDEntity, ButtonEntity):
@wled_exception_handler @wled_exception_handler
async def async_press(self) -> None: async def async_press(self) -> None:
"""Send out a restart command.""" """Send out a update command."""
current = self.coordinator.data.info.version current = self.coordinator.data.info.version
beta = self.coordinator.data.info.version_latest_beta beta = self.coordinator.data.info.version_latest_beta
stable = self.coordinator.data.info.version_latest_stable stable = self.coordinator.data.info.version_latest_stable
# If we already run a pre-release, allow upgrading to a newer # If we already run a pre-release, allow update to a newer
# pre-release or newer stable, otherwise, offer a normal stable upgrades. # pre-release or newer stable, otherwise, offer a normal stable updates.
version = stable version = stable
if ( if (
current is not None current is not None

View File

@ -5,10 +5,14 @@ from freezegun import freeze_time
import pytest import pytest
from wled import WLEDConnectionError, WLEDError from wled import WLEDConnectionError, WLEDError
from homeassistant.components.button import DOMAIN as BUTTON_DOMAIN, SERVICE_PRESS from homeassistant.components.button import (
DOMAIN as BUTTON_DOMAIN,
SERVICE_PRESS,
ButtonDeviceClass,
)
from homeassistant.const import ( from homeassistant.const import (
ATTR_DEVICE_CLASS,
ATTR_ENTITY_ID, ATTR_ENTITY_ID,
ATTR_ICON,
ENTITY_CATEGORY_CONFIG, ENTITY_CATEGORY_CONFIG,
STATE_UNAVAILABLE, STATE_UNAVAILABLE,
STATE_UNKNOWN, STATE_UNKNOWN,
@ -27,8 +31,8 @@ async def test_button_restart(
state = hass.states.get("button.wled_rgb_light_restart") state = hass.states.get("button.wled_rgb_light_restart")
assert state assert state
assert state.attributes.get(ATTR_ICON) == "mdi:restart"
assert state.state == STATE_UNKNOWN assert state.state == STATE_UNKNOWN
assert state.attributes[ATTR_DEVICE_CLASS] == ButtonDeviceClass.RESTART
entry = entity_registry.async_get("button.wled_rgb_light_restart") entry = entity_registry.async_get("button.wled_rgb_light_restart")
assert entry assert entry
@ -93,31 +97,31 @@ async def test_button_connection_error(
assert "Error communicating with API" in caplog.text assert "Error communicating with API" in caplog.text
async def test_button_upgrade_stay_stable( async def test_button_update_stay_stable(
hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
) -> None: ) -> None:
"""Test the upgrade button. """Test the update button.
There is both an upgrade for beta and stable available, however, the device There is both an update for beta and stable available, however, the device
is currently running a stable version. Therefore, the upgrade button should is currently running a stable version. Therefore, the update button should
upgrade the the next stable (even though beta is newer). update the the next stable (even though beta is newer).
""" """
entity_registry = er.async_get(hass) entity_registry = er.async_get(hass)
entry = entity_registry.async_get("button.wled_rgb_light_upgrade") entry = entity_registry.async_get("button.wled_rgb_light_update")
assert entry assert entry
assert entry.unique_id == "aabbccddeeff_upgrade" assert entry.unique_id == "aabbccddeeff_update"
assert entry.entity_category == ENTITY_CATEGORY_CONFIG assert entry.entity_category == ENTITY_CATEGORY_CONFIG
state = hass.states.get("button.wled_rgb_light_upgrade") state = hass.states.get("button.wled_rgb_light_update")
assert state assert state
assert state.attributes.get(ATTR_ICON) == "mdi:cellphone-arrow-down"
assert state.state == STATE_UNKNOWN assert state.state == STATE_UNKNOWN
assert state.attributes[ATTR_DEVICE_CLASS] == ButtonDeviceClass.UPDATE
await hass.services.async_call( await hass.services.async_call(
BUTTON_DOMAIN, BUTTON_DOMAIN,
SERVICE_PRESS, SERVICE_PRESS,
{ATTR_ENTITY_ID: "button.wled_rgb_light_upgrade"}, {ATTR_ENTITY_ID: "button.wled_rgb_light_update"},
blocking=True, blocking=True,
) )
await hass.async_block_till_done() await hass.async_block_till_done()
@ -126,19 +130,19 @@ async def test_button_upgrade_stay_stable(
@pytest.mark.parametrize("mock_wled", ["wled/rgbw.json"], indirect=True) @pytest.mark.parametrize("mock_wled", ["wled/rgbw.json"], indirect=True)
async def test_button_upgrade_beta_to_stable( async def test_button_update_beta_to_stable(
hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
) -> None: ) -> None:
"""Test the upgrade button. """Test the update button.
There is both an upgrade for beta and stable available the device There is both an update for beta and stable available the device
is currently a beta, however, a newer stable is available. Therefore, the is currently a beta, however, a newer stable is available. Therefore, the
upgrade button should upgrade to the next stable. update button should update to the next stable.
""" """
await hass.services.async_call( await hass.services.async_call(
BUTTON_DOMAIN, BUTTON_DOMAIN,
SERVICE_PRESS, SERVICE_PRESS,
{ATTR_ENTITY_ID: "button.wled_rgbw_light_upgrade"}, {ATTR_ENTITY_ID: "button.wled_rgbw_light_update"},
blocking=True, blocking=True,
) )
await hass.async_block_till_done() await hass.async_block_till_done()
@ -147,18 +151,18 @@ async def test_button_upgrade_beta_to_stable(
@pytest.mark.parametrize("mock_wled", ["wled/rgb_single_segment.json"], indirect=True) @pytest.mark.parametrize("mock_wled", ["wled/rgb_single_segment.json"], indirect=True)
async def test_button_upgrade_stay_beta( async def test_button_update_stay_beta(
hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
) -> None: ) -> None:
"""Test the upgrade button. """Test the update button.
There is an upgrade for beta and the device is currently a beta. Therefore, There is an update for beta and the device is currently a beta. Therefore,
the upgrade button should upgrade to the next beta. the update button should update to the next beta.
""" """
await hass.services.async_call( await hass.services.async_call(
BUTTON_DOMAIN, BUTTON_DOMAIN,
SERVICE_PRESS, SERVICE_PRESS,
{ATTR_ENTITY_ID: "button.wled_rgb_light_upgrade"}, {ATTR_ENTITY_ID: "button.wled_rgb_light_update"},
blocking=True, blocking=True,
) )
await hass.async_block_till_done() await hass.async_block_till_done()
@ -167,10 +171,10 @@ async def test_button_upgrade_stay_beta(
@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True) @pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True)
async def test_button_no_upgrade_available( async def test_button_no_update_available(
hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock
) -> None: ) -> None:
"""Test the upgrade button. There is no update available.""" """Test the update button. There is no update available."""
state = hass.states.get("button.wled_websocket_upgrade") state = hass.states.get("button.wled_websocket_update")
assert state assert state
assert state.state == STATE_UNAVAILABLE assert state.state == STATE_UNAVAILABLE