From 0950ab0dd8ceb0fc6030a1032cf26e4d1585f2e7 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Wed, 3 Jun 2020 17:18:50 +0200 Subject: [PATCH] Fix dynamically add/remove WLED strip segments (#36407) --- homeassistant/components/wled/light.py | 66 ++++++- tests/components/wled/test_light.py | 36 +++- tests/fixtures/wled/rgb_single_segment.json | 202 ++++++++++++++++++++ 3 files changed, 298 insertions(+), 6 deletions(-) create mode 100644 tests/fixtures/wled/rgb_single_segment.json diff --git a/homeassistant/components/wled/light.py b/homeassistant/components/wled/light.py index 6b86be6265b..77f8960ada2 100644 --- a/homeassistant/components/wled/light.py +++ b/homeassistant/components/wled/light.py @@ -1,4 +1,5 @@ """Support for LED lights.""" +from functools import partial import logging from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -20,8 +21,12 @@ from homeassistant.components.light import ( LightEntity, ) from homeassistant.config_entries import ConfigEntry +from homeassistant.core import callback from homeassistant.helpers import config_validation as cv, entity_platform from homeassistant.helpers.entity import Entity +from homeassistant.helpers.entity_registry import ( + async_get_registry as async_get_entity_registry, +) from homeassistant.helpers.typing import HomeAssistantType import homeassistant.util.color as color_util @@ -70,12 +75,12 @@ async def async_setup_entry( "async_effect", ) - lights = [ - WLEDLight(entry.entry_id, coordinator, light.segment_id) - for light in coordinator.data.state.segments - ] + update_segments = partial( + async_update_segments, entry, coordinator, {}, async_add_entities + ) - async_add_entities(lights, True) + coordinator.async_add_listener(update_segments) + update_segments() class WLEDLight(LightEntity, WLEDDeviceEntity): @@ -105,6 +110,16 @@ class WLEDLight(LightEntity, WLEDDeviceEntity): """Return the unique ID for this sensor.""" return f"{self.coordinator.data.info.mac_address}_{self._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 device_state_attributes(self) -> Optional[Dict[str, Any]]: """Return the state attributes of the entity.""" @@ -259,3 +274,44 @@ class WLEDLight(LightEntity, WLEDDeviceEntity): data[ATTR_SPEED] = speed await self.coordinator.wled.light(**data) + + +@callback +def async_update_segments( + entry: ConfigEntry, + coordinator: WLEDDataUpdateCoordinator, + current: Dict[int, WLEDLight], + async_add_entities, +) -> None: + """Update segments.""" + segment_ids = {light.segment_id for light in coordinator.data.state.segments} + current_ids = set(current) + + # Process new segments, add them to Home Assistant + new_segments = [] + for segment_id in segment_ids - current_ids: + current[segment_id] = WLEDLight(entry.entry_id, coordinator, segment_id) + new_segments.append(current[segment_id]) + + if new_segments: + async_add_entities(new_segments) + + # Process deleted segments, remove them from Home Assistant + for segment_id in current_ids - segment_ids: + coordinator.hass.async_create_task( + async_remove_segment(segment_id, coordinator, current) + ) + + +async def async_remove_segment( + segment_id: int, + coordinator: WLEDDataUpdateCoordinator, + current: Dict[int, WLEDLight], +) -> None: + """Remove WLED segment light from Home Assistant.""" + entity = current[segment_id] + await entity.async_remove() + registry = await async_get_entity_registry(coordinator.hass) + if entity.entity_id in registry.entities: + registry.async_remove(entity.entity_id) + del current[segment_id] diff --git a/tests/components/wled/test_light.py b/tests/components/wled/test_light.py index 307f1e56411..8854b00ff83 100644 --- a/tests/components/wled/test_light.py +++ b/tests/components/wled/test_light.py @@ -1,5 +1,7 @@ """Tests for the WLED light platform.""" -from wled import WLEDConnectionError +import json + +from wled import Device as WLEDDevice, WLEDConnectionError from homeassistant.components.light import ( ATTR_BRIGHTNESS, @@ -11,6 +13,7 @@ from homeassistant.components.light import ( ATTR_WHITE_VALUE, DOMAIN as LIGHT_DOMAIN, ) +from homeassistant.components.wled import SCAN_INTERVAL from homeassistant.components.wled.const import ( ATTR_INTENSITY, ATTR_PALETTE, @@ -30,8 +33,10 @@ from homeassistant.const import ( STATE_UNAVAILABLE, ) from homeassistant.core import HomeAssistant +import homeassistant.util.dt as dt_util from tests.async_mock import patch +from tests.common import async_fire_time_changed, load_fixture from tests.components.wled import init_integration from tests.test_util.aiohttp import AiohttpClientMocker @@ -137,6 +142,35 @@ async def test_switch_change_state( ) +async def test_dynamically_handle_segments( + hass: HomeAssistant, aioclient_mock: AiohttpClientMocker +) -> None: + """Test if a new/deleted segment is dynamically added/removed.""" + await init_integration(hass, aioclient_mock) + + assert hass.states.get("light.wled_rgb_light") + assert hass.states.get("light.wled_rgb_light_1") + + data = json.loads(load_fixture("wled/rgb_single_segment.json")) + device = WLEDDevice(data) + + # Test removal if segment went missing + with patch( + "homeassistant.components.wled.WLED.update", return_value=device, + ): + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + assert hass.states.get("light.wled_rgb_light") + assert not hass.states.get("light.wled_rgb_light_1") + + # Test adding if segment shows up again + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + + assert hass.states.get("light.wled_rgb_light") + assert hass.states.get("light.wled_rgb_light_1") + + async def test_light_error( hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, caplog ) -> None: diff --git a/tests/fixtures/wled/rgb_single_segment.json b/tests/fixtures/wled/rgb_single_segment.json new file mode 100644 index 00000000000..e53ce680ece --- /dev/null +++ b/tests/fixtures/wled/rgb_single_segment.json @@ -0,0 +1,202 @@ +{ + "state": { + "on": true, + "bri": 127, + "transition": 7, + "ps": -1, + "pl": -1, + "nl": { + "on": false, + "dur": 60, + "fade": true, + "tbri": 0 + }, + "udpn": { + "send": false, + "recv": true + }, + "seg": [ + { + "id": 0, + "start": 0, + "stop": 30, + "len": 20, + "col": [[255, 159, 0], [0, 0, 0], [0, 0, 0]], + "fx": 0, + "sx": 32, + "ix": 128, + "pal": 0, + "sel": true, + "rev": false, + "cln": -1 + } + ] + }, + "info": { + "ver": "0.8.5", + "vid": 1909122, + "leds": { + "count": 30, + "rgbw": false, + "pin": [2], + "pwr": 470, + "maxpwr": 850, + "maxseg": 10 + }, + "name": "WLED RGB Light", + "udpport": 21324, + "live": false, + "fxcount": 81, + "palcount": 50, + "wifi": { + "bssid": "AA:AA:AA:AA:AA:BB", + "rssi": -62, + "signal": 76, + "channel": 11 + }, + "arch": "esp8266", + "core": "2_4_2", + "freeheap": 14600, + "uptime": 32, + "opt": 119, + "brand": "WLED", + "product": "DIY light", + "btype": "bin", + "mac": "aabbccddeeff" + }, + "effects": [ + "Solid", + "Blink", + "Breathe", + "Wipe", + "Wipe Random", + "Random Colors", + "Sweep", + "Dynamic", + "Colorloop", + "Rainbow", + "Scan", + "Dual Scan", + "Fade", + "Chase", + "Chase Rainbow", + "Running", + "Saw", + "Twinkle", + "Dissolve", + "Dissolve Rnd", + "Sparkle", + "Dark Sparkle", + "Sparkle+", + "Strobe", + "Strobe Rainbow", + "Mega Strobe", + "Blink Rainbow", + "Android", + "Chase", + "Chase Random", + "Chase Rainbow", + "Chase Flash", + "Chase Flash Rnd", + "Rainbow Runner", + "Colorful", + "Traffic Light", + "Sweep Random", + "Running 2", + "Red & Blue", + "Stream", + "Scanner", + "Lighthouse", + "Fireworks", + "Rain", + "Merry Christmas", + "Fire Flicker", + "Gradient", + "Loading", + "In Out", + "In In", + "Out Out", + "Out In", + "Circus", + "Halloween", + "Tri Chase", + "Tri Wipe", + "Tri Fade", + "Lightning", + "ICU", + "Multi Comet", + "Dual Scanner", + "Stream 2", + "Oscillate", + "Pride 2015", + "Juggle", + "Palette", + "Fire 2012", + "Colorwaves", + "BPM", + "Fill Noise", + "Noise 1", + "Noise 2", + "Noise 3", + "Noise 4", + "Colortwinkle", + "Lake", + "Meteor", + "Smooth Meteor", + "Railway", + "Ripple", + "Twinklefox" + ], + "palettes": [ + "Default", + "Random Cycle", + "Primary Color", + "Based on Primary", + "Set Colors", + "Based on Set", + "Party", + "Cloud", + "Lava", + "Ocean", + "Forest", + "Rainbow", + "Rainbow Bands", + "Sunset", + "Rivendell", + "Breeze", + "Red & Blue", + "Yellowout", + "Analogous", + "Splash", + "Pastel", + "Sunset 2", + "Beech", + "Vintage", + "Departure", + "Landscape", + "Beach", + "Sherbet", + "Hult", + "Hult 64", + "Drywet", + "Jul", + "Grintage", + "Rewhi", + "Tertiary", + "Fire", + "Icefire", + "Cyane", + "Light Pink", + "Autumn", + "Magenta", + "Magred", + "Yelmag", + "Yelblu", + "Orange & Teal", + "Tiamat", + "April Night", + "Orangery", + "C9", + "Sakura" + ] +}