From 2f346a8048bc2d00ed17ef1be59d35452e1dd4b2 Mon Sep 17 00:00:00 2001 From: Teemu R Date: Tue, 26 Oct 2021 09:04:19 +0200 Subject: [PATCH] Add speed & intensity controls to wled (#56862) Co-authored-by: Franck Nijhof --- homeassistant/components/wled/__init__.py | 3 +- homeassistant/components/wled/number.py | 133 ++++++++++ tests/components/wled/test_number.py | 302 ++++++++++++++++++++++ 3 files changed, 437 insertions(+), 1 deletion(-) create mode 100644 homeassistant/components/wled/number.py create mode 100644 tests/components/wled/test_number.py diff --git a/homeassistant/components/wled/__init__.py b/homeassistant/components/wled/__init__.py index 29c6b98b381..e7697676014 100644 --- a/homeassistant/components/wled/__init__.py +++ b/homeassistant/components/wled/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations 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 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN @@ -11,7 +12,7 @@ from homeassistant.core import HomeAssistant from .const import DOMAIN from .coordinator import WLEDDataUpdateCoordinator -PLATFORMS = (LIGHT_DOMAIN, SELECT_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN) +PLATFORMS = (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/number.py b/homeassistant/components/wled/number.py new file mode 100644 index 00000000000..a51eb915ba6 --- /dev/null +++ b/homeassistant/components/wled/number.py @@ -0,0 +1,133 @@ +"""Support for LED numbers.""" +from __future__ import annotations + +from functools import partial + +from homeassistant.components.number import NumberEntity, NumberEntityDescription +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import ENTITY_CATEGORY_CONFIG +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.entity_platform import AddEntitiesCallback + +from .const import ATTR_INTENSITY, ATTR_SPEED, DOMAIN +from .coordinator import WLEDDataUpdateCoordinator +from .helpers import wled_exception_handler +from .models import WLEDEntity + +PARALLEL_UPDATES = 1 + + +async def async_setup_entry( + hass: HomeAssistant, + entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +) -> None: + """Set up WLED number based on a config entry.""" + coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + + update_segments = partial( + async_update_segments, + coordinator, + set(), + async_add_entities, + ) + coordinator.async_add_listener(update_segments) + update_segments() + + +NUMBERS = [ + NumberEntityDescription( + key=ATTR_SPEED, + name="Speed", + icon="mdi:speedometer", + entity_category=ENTITY_CATEGORY_CONFIG, + ), + NumberEntityDescription( + key=ATTR_INTENSITY, + name="Intensity", + entity_category=ENTITY_CATEGORY_CONFIG, + ), +] + + +class WLEDNumber(WLEDEntity, NumberEntity): + """Defines a WLED speed number.""" + + _attr_step = 1 + _attr_min_value = 0 + _attr_max_value = 255 + + def __init__( + self, + coordinator: WLEDDataUpdateCoordinator, + segment: int, + description: NumberEntityDescription, + ) -> None: + """Initialize WLED .""" + super().__init__(coordinator=coordinator) + self.entity_description = description + + # Segment 0 uses a simpler name, which is more natural for when using + # a single segment / using WLED with one big LED strip. + self._attr_name = ( + f"{coordinator.data.info.name} Segment {segment} {description.name}" + ) + if segment == 0: + self._attr_name = f"{coordinator.data.info.name} {description.name}" + + self._attr_unique_id = ( + f"{coordinator.data.info.mac_address}_{description.key}_{segment}" + ) + self._segment = segment + + @property + def available(self) -> bool: + """Return True if entity is available.""" + try: + self.coordinator.data.state.segments[self._segment] + except IndexError: + return False + + return super().available + + @property + def value(self) -> float | None: + """Return the current WLED segment number value.""" + return getattr( + self.coordinator.data.state.segments[self._segment], + self.entity_description.key, + ) + + @wled_exception_handler + async def async_set_value(self, value: float) -> None: + """Set the WLED segment value.""" + key = self.entity_description.key + if key == ATTR_SPEED: + await self.coordinator.wled.segment( + segment_id=self._segment, speed=int(value) + ) + elif key == ATTR_INTENSITY: + return await self.coordinator.wled.segment( + segment_id=self._segment, intensity=int(value) + ) + + +@callback +def async_update_segments( + coordinator: WLEDDataUpdateCoordinator, + current_ids: set[int], + async_add_entities, +) -> None: + """Update segments.""" + segment_ids = {segment.segment_id for segment in coordinator.data.state.segments} + + new_entities = [] + + # Process new segments, add them to Home Assistant + for segment_id in segment_ids - current_ids: + current_ids.add(segment_id) + for desc in NUMBERS: + new_entities.append(WLEDNumber(coordinator, segment_id, desc)) + + if new_entities: + async_add_entities(new_entities) diff --git a/tests/components/wled/test_number.py b/tests/components/wled/test_number.py new file mode 100644 index 00000000000..e4b8958b077 --- /dev/null +++ b/tests/components/wled/test_number.py @@ -0,0 +1,302 @@ +"""Tests for the WLED number platform.""" +import json +from unittest.mock import MagicMock + +import pytest +from wled import Device as WLEDDevice, WLEDConnectionError, WLEDError + +from homeassistant.components.number import ATTR_MAX, ATTR_MIN, DOMAIN as NUMBER_DOMAIN +from homeassistant.components.number.const import ( + ATTR_STEP, + ATTR_VALUE, + SERVICE_SET_VALUE, +) +from homeassistant.components.wled.const import SCAN_INTERVAL +from homeassistant.const import ATTR_ENTITY_ID, ATTR_ICON, STATE_UNAVAILABLE +from homeassistant.core import HomeAssistant +from homeassistant.helpers import entity_registry as er +import homeassistant.util.dt as dt_util + +from tests.common import MockConfigEntry, async_fire_time_changed, load_fixture + + +async def test_speed_state( + hass: HomeAssistant, init_integration: MockConfigEntry +) -> None: + """Test the creation and values of the WLED numbers.""" + entity_registry = er.async_get(hass) + + # First segment of the strip + state = hass.states.get("number.wled_rgb_light_segment_1_speed") + assert state + assert state.attributes.get(ATTR_ICON) == "mdi:speedometer" + assert state.attributes.get(ATTR_MAX) == 255 + assert state.attributes.get(ATTR_MIN) == 0 + assert state.attributes.get(ATTR_STEP) == 1 + assert state.state == "16" + + entry = entity_registry.async_get("number.wled_rgb_light_segment_1_speed") + assert entry + assert entry.unique_id == "aabbccddeeff_speed_1" + + +async def test_speed_segment_change_state( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Test the value change of the WLED segments.""" + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + { + ATTR_ENTITY_ID: "number.wled_rgb_light_segment_1_speed", + ATTR_VALUE: 42, + }, + blocking=True, + ) + await hass.async_block_till_done() + assert mock_wled.segment.call_count == 1 + mock_wled.segment.assert_called_with( + segment_id=1, + speed=42, + ) + + +@pytest.mark.parametrize("mock_wled", ["wled/rgb_single_segment.json"], indirect=True) +async def test_speed_dynamically_handle_segments( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Test if a new/deleted segment is dynamically added/removed.""" + segment0 = hass.states.get("number.wled_rgb_light_speed") + segment1 = hass.states.get("number.wled_rgb_light_segment_1_speed") + assert segment0 + assert segment0.state == "32" + assert not segment1 + + # Test adding a segment dynamically... + return_value = mock_wled.update.return_value + mock_wled.update.return_value = WLEDDevice( + json.loads(load_fixture("wled/rgb.json")) + ) + + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + + segment0 = hass.states.get("number.wled_rgb_light_speed") + segment1 = hass.states.get("number.wled_rgb_light_segment_1_speed") + assert segment0 + assert segment0.state == "32" + assert segment1 + assert segment1.state == "16" + + # Test remove segment again... + mock_wled.update.return_value = return_value + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + + segment0 = hass.states.get("number.wled_rgb_light_speed") + segment1 = hass.states.get("number.wled_rgb_light_segment_1_speed") + assert segment0 + assert segment0.state == "32" + assert segment1 + assert segment1.state == STATE_UNAVAILABLE + + +async def test_speed_error( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test error handling of the WLED numbers.""" + mock_wled.segment.side_effect = WLEDError + + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + { + ATTR_ENTITY_ID: "number.wled_rgb_light_segment_1_speed", + ATTR_VALUE: 42, + }, + blocking=True, + ) + await hass.async_block_till_done() + + state = hass.states.get("number.wled_rgb_light_segment_1_speed") + assert state + assert state.state == "16" + assert "Invalid response from API" in caplog.text + assert mock_wled.segment.call_count == 1 + mock_wled.segment.assert_called_with(segment_id=1, speed=42) + + +async def test_speed_connection_error( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test error handling of the WLED numbers.""" + mock_wled.segment.side_effect = WLEDConnectionError + + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + { + ATTR_ENTITY_ID: "number.wled_rgb_light_segment_1_speed", + ATTR_VALUE: 42, + }, + blocking=True, + ) + await hass.async_block_till_done() + + state = hass.states.get("number.wled_rgb_light_segment_1_speed") + assert state + assert state.state == STATE_UNAVAILABLE + assert "Error communicating with API" in caplog.text + assert mock_wled.segment.call_count == 1 + mock_wled.segment.assert_called_with(segment_id=1, speed=42) + + +async def test_intensity_state( + hass: HomeAssistant, init_integration: MockConfigEntry +) -> None: + """Test the creation and values of the WLED numbers.""" + entity_registry = er.async_get(hass) + + # First segment of the strip + state = hass.states.get("number.wled_rgb_light_segment_1_intensity") + assert state + assert state.attributes.get(ATTR_ICON) is None + assert state.attributes.get(ATTR_MAX) == 255 + assert state.attributes.get(ATTR_MIN) == 0 + assert state.attributes.get(ATTR_STEP) == 1 + assert state.state == "64" + + entry = entity_registry.async_get("number.wled_rgb_light_segment_1_intensity") + assert entry + assert entry.unique_id == "aabbccddeeff_intensity_1" + + +async def test_intensity_segment_change_state( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Test the value change of the WLED segments.""" + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + { + ATTR_ENTITY_ID: "number.wled_rgb_light_segment_1_intensity", + ATTR_VALUE: 128, + }, + blocking=True, + ) + await hass.async_block_till_done() + assert mock_wled.segment.call_count == 1 + mock_wled.segment.assert_called_with( + segment_id=1, + intensity=128, + ) + + +@pytest.mark.parametrize("mock_wled", ["wled/rgb_single_segment.json"], indirect=True) +async def test_intensity_dynamically_handle_segments( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Test if a new/deleted segment is dynamically added/removed.""" + segment0 = hass.states.get("number.wled_rgb_light_intensity") + segment1 = hass.states.get("number.wled_rgb_light_segment_1_intensity") + assert segment0 + assert segment0.state == "128" + assert not segment1 + + # Test adding a segment dynamically... + return_value = mock_wled.update.return_value + mock_wled.update.return_value = WLEDDevice( + json.loads(load_fixture("wled/rgb.json")) + ) + + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + + segment0 = hass.states.get("number.wled_rgb_light_intensity") + segment1 = hass.states.get("number.wled_rgb_light_segment_1_intensity") + assert segment0 + assert segment0.state == "128" + assert segment1 + assert segment1.state == "64" + + # Test remove segment again... + mock_wled.update.return_value = return_value + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + + segment0 = hass.states.get("number.wled_rgb_light_intensity") + segment1 = hass.states.get("number.wled_rgb_light_segment_1_intensity") + assert segment0 + assert segment0.state == "128" + assert segment1 + assert segment1.state == STATE_UNAVAILABLE + + +async def test_intensity_error( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test error handling of the WLED numbers.""" + mock_wled.segment.side_effect = WLEDError + + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + { + ATTR_ENTITY_ID: "number.wled_rgb_light_segment_1_intensity", + ATTR_VALUE: 21, + }, + blocking=True, + ) + await hass.async_block_till_done() + + state = hass.states.get("number.wled_rgb_light_segment_1_intensity") + assert state + assert state.state == "64" + assert "Invalid response from API" in caplog.text + assert mock_wled.segment.call_count == 1 + mock_wled.segment.assert_called_with(segment_id=1, intensity=21) + + +async def test_intensity_connection_error( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test error handling of the WLED numbers.""" + mock_wled.segment.side_effect = WLEDConnectionError + + await hass.services.async_call( + NUMBER_DOMAIN, + SERVICE_SET_VALUE, + { + ATTR_ENTITY_ID: "number.wled_rgb_light_segment_1_intensity", + ATTR_VALUE: 128, + }, + blocking=True, + ) + await hass.async_block_till_done() + + state = hass.states.get("number.wled_rgb_light_segment_1_intensity") + assert state + assert state.state == STATE_UNAVAILABLE + assert "Error communicating with API" in caplog.text + assert mock_wled.segment.call_count == 1 + mock_wled.segment.assert_called_with(segment_id=1, intensity=128)