From b83b82ca7d04827b4b2cd1210ae023732736ca2e Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Fri, 11 Jun 2021 20:55:08 +0200 Subject: [PATCH] WLED WebSocket support - local push updates (#51683) Co-authored-by: Martin Hjelmare Co-authored-by: Paulus Schoutsen --- homeassistant/components/wled/__init__.py | 8 + homeassistant/components/wled/const.py | 2 +- homeassistant/components/wled/coordinator.py | 65 ++++- homeassistant/components/wled/manifest.json | 2 +- homeassistant/components/wled/sensor.py | 2 +- homeassistant/components/wled/switch.py | 2 +- tests/components/wled/conftest.py | 1 + tests/components/wled/test_coordinator.py | 196 +++++++++++++ tests/components/wled/test_init.py | 20 ++ tests/fixtures/wled/rgb_websocket.json | 289 +++++++++++++++++++ 10 files changed, 580 insertions(+), 7 deletions(-) create mode 100644 tests/components/wled/test_coordinator.py create mode 100644 tests/fixtures/wled/rgb_websocket.json diff --git a/homeassistant/components/wled/__init__.py b/homeassistant/components/wled/__init__.py index bd8316a2ff0..2082f45c6c3 100644 --- a/homeassistant/components/wled/__init__.py +++ b/homeassistant/components/wled/__init__.py @@ -38,5 +38,13 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload WLED config entry.""" unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) if unload_ok: + coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + + # Ensure disconnected and cleanup stop sub + await coordinator.wled.disconnect() + if coordinator.unsub: + coordinator.unsub() + del hass.data[DOMAIN][entry.entry_id] + return unload_ok diff --git a/homeassistant/components/wled/const.py b/homeassistant/components/wled/const.py index 8f759ea3e90..77e404cdd4c 100644 --- a/homeassistant/components/wled/const.py +++ b/homeassistant/components/wled/const.py @@ -6,7 +6,7 @@ import logging DOMAIN = "wled" LOGGER = logging.getLogger(__package__) -SCAN_INTERVAL = timedelta(seconds=5) +SCAN_INTERVAL = timedelta(seconds=10) # Attributes ATTR_COLOR_PRIMARY = "color_primary" diff --git a/homeassistant/components/wled/coordinator.py b/homeassistant/components/wled/coordinator.py index 8a9312dfc0a..16d56705879 100644 --- a/homeassistant/components/wled/coordinator.py +++ b/homeassistant/components/wled/coordinator.py @@ -1,8 +1,13 @@ """DataUpdateCoordinator for WLED.""" +from __future__ import annotations -from wled import WLED, Device as WLEDDevice, WLEDError +import asyncio +from typing import Callable -from homeassistant.core import HomeAssistant +from wled import WLED, Device as WLEDDevice, WLEDConnectionClosed, WLEDError + +from homeassistant.const import EVENT_HOMEASSISTANT_STOP +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed @@ -20,6 +25,7 @@ class WLEDDataUpdateCoordinator(DataUpdateCoordinator[WLEDDevice]): ) -> None: """Initialize global WLED data updater.""" self.wled = WLED(host, session=async_get_clientsession(hass)) + self.unsub: Callable | None = None super().__init__( hass, @@ -33,9 +39,62 @@ class WLEDDataUpdateCoordinator(DataUpdateCoordinator[WLEDDevice]): for update_callback in self._listeners: update_callback() + @callback + def _use_websocket(self) -> None: + """Use WebSocket for updates, instead of polling.""" + + async def listen() -> None: + """Listen for state changes via WebSocket.""" + try: + await self.wled.connect() + except WLEDError as err: + self.logger.info(err) + if self.unsub: + self.unsub() + self.unsub = None + return + + try: + await self.wled.listen(callback=self.async_set_updated_data) + except WLEDConnectionClosed as err: + self.last_update_success = False + self.logger.info(err) + except WLEDError as err: + self.last_update_success = False + self.update_listeners() + self.logger.error(err) + + # Ensure we are disconnected + await self.wled.disconnect() + if self.unsub: + self.unsub() + self.unsub = None + + async def close_websocket(_) -> None: + """Close WebSocket connection.""" + await self.wled.disconnect() + + # Clean disconnect WebSocket on Home Assistant shutdown + self.unsub = self.hass.bus.async_listen_once( + EVENT_HOMEASSISTANT_STOP, close_websocket + ) + + # Start listening + asyncio.create_task(listen()) + async def _async_update_data(self) -> WLEDDevice: """Fetch data from WLED.""" try: - return await self.wled.update(full_update=not self.last_update_success) + device = await self.wled.update(full_update=not self.last_update_success) except WLEDError as error: raise UpdateFailed(f"Invalid response from API: {error}") from error + + # If the device supports a WebSocket, try activating it. + if ( + device.info.websocket is not None + and not self.wled.connected + and not self.unsub + ): + self._use_websocket() + + return device diff --git a/homeassistant/components/wled/manifest.json b/homeassistant/components/wled/manifest.json index b5ac91a1bf8..91a449e918d 100644 --- a/homeassistant/components/wled/manifest.json +++ b/homeassistant/components/wled/manifest.json @@ -7,5 +7,5 @@ "zeroconf": ["_wled._tcp.local."], "codeowners": ["@frenck"], "quality_scale": "platinum", - "iot_class": "local_polling" + "iot_class": "local_push" } diff --git a/homeassistant/components/wled/sensor.py b/homeassistant/components/wled/sensor.py index 4008c42f292..37311e333c3 100644 --- a/homeassistant/components/wled/sensor.py +++ b/homeassistant/components/wled/sensor.py @@ -40,7 +40,7 @@ async def async_setup_entry( WLEDWifiSignalSensor(coordinator), ] - async_add_entities(sensors, True) + async_add_entities(sensors) class WLEDEstimatedCurrentSensor(WLEDEntity, SensorEntity): diff --git a/homeassistant/components/wled/switch.py b/homeassistant/components/wled/switch.py index 74f58472a19..b17572f7607 100644 --- a/homeassistant/components/wled/switch.py +++ b/homeassistant/components/wled/switch.py @@ -35,7 +35,7 @@ async def async_setup_entry( WLEDSyncSendSwitch(coordinator), WLEDSyncReceiveSwitch(coordinator), ] - async_add_entities(switches, True) + async_add_entities(switches) class WLEDNightlightSwitch(WLEDEntity, SwitchEntity): diff --git a/tests/components/wled/conftest.py b/tests/components/wled/conftest.py index f66171b6025..80b351a20f1 100644 --- a/tests/components/wled/conftest.py +++ b/tests/components/wled/conftest.py @@ -63,6 +63,7 @@ def mock_wled(request: pytest.FixtureRequest) -> Generator[None, MagicMock, None ) as wled_mock: wled = wled_mock.return_value wled.update.return_value = device + wled.connected = False yield wled diff --git a/tests/components/wled/test_coordinator.py b/tests/components/wled/test_coordinator.py new file mode 100644 index 00000000000..47190604238 --- /dev/null +++ b/tests/components/wled/test_coordinator.py @@ -0,0 +1,196 @@ +"""Tests for the coordinator of the WLED integration.""" +import asyncio +from copy import deepcopy +from typing import Callable +from unittest.mock import MagicMock + +import pytest +from wled import ( + Device as WLEDDevice, + WLEDConnectionClosed, + WLEDConnectionError, + WLEDError, +) + +from homeassistant.components.wled.const import SCAN_INTERVAL +from homeassistant.const import ( + EVENT_HOMEASSISTANT_STOP, + STATE_OFF, + STATE_ON, + STATE_UNAVAILABLE, +) +from homeassistant.core import HomeAssistant +import homeassistant.util.dt as dt_util + +from tests.common import MockConfigEntry, async_fire_time_changed + + +async def test_not_supporting_websocket( + hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock +) -> None: + """Ensure no WebSocket attempt is made if non-WebSocket device.""" + assert mock_wled.connect.call_count == 0 + + +@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True) +async def test_websocket_already_connected( + hass: HomeAssistant, init_integration: MockConfigEntry, mock_wled: MagicMock +) -> None: + """Ensure no a second WebSocket connection is made, if already connected.""" + assert mock_wled.connect.call_count == 1 + + mock_wled.connected = True + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + + assert mock_wled.connect.call_count == 1 + + +@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True) +async def test_websocket_connect_error_no_listen( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Ensure we don't start listening if WebSocket connection failed.""" + assert mock_wled.connect.call_count == 1 + assert mock_wled.listen.call_count == 1 + + mock_wled.connect.side_effect = WLEDConnectionError + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await hass.async_block_till_done() + + assert mock_wled.connect.call_count == 2 + assert mock_wled.listen.call_count == 1 + + +@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True) +async def test_websocket( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Test WebSocket connection.""" + state = hass.states.get("light.wled_websocket") + assert state + assert state.state == STATE_ON + + # There is no Future in place yet... + assert mock_wled.connect.call_count == 1 + assert mock_wled.listen.call_count == 1 + assert mock_wled.disconnect.call_count == 1 + + connection_connected = asyncio.Future() + connection_finished = asyncio.Future() + + async def connect(callback: Callable[[WLEDDevice], None]): + connection_connected.set_result(callback) + await connection_finished + + # Mock out wled.listen with a Future + mock_wled.listen.side_effect = connect + + # Mock out the event bus + mock_bus = MagicMock() + hass.bus = mock_bus + + # Next refresh it should connect + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + callback = await connection_connected + + # Connected to WebSocket, disconnect not called + # listening for Home Assistant to stop + assert mock_wled.connect.call_count == 2 + assert mock_wled.listen.call_count == 2 + assert mock_wled.disconnect.call_count == 1 + assert mock_bus.async_listen_once.call_count == 1 + assert ( + mock_bus.async_listen_once.call_args_list[0][0][0] == EVENT_HOMEASSISTANT_STOP + ) + assert ( + mock_bus.async_listen_once.call_args_list[0][0][1].__name__ == "close_websocket" + ) + assert mock_bus.async_listen_once.return_value.call_count == 0 + + # Send update from WebSocket + updated_device = deepcopy(mock_wled.update.return_value) + updated_device.state.on = False + callback(updated_device) + await hass.async_block_till_done() + + # Check if entity updated + state = hass.states.get("light.wled_websocket") + assert state + assert state.state == STATE_OFF + + # Resolve Future with a connection losed. + connection_finished.set_exception(WLEDConnectionClosed) + await hass.async_block_till_done() + + # Disconnect called, unsubbed Home Assistant stop listener + assert mock_wled.disconnect.call_count == 2 + assert mock_bus.async_listen_once.return_value.call_count == 1 + + # Light still available, as polling takes over + state = hass.states.get("light.wled_websocket") + assert state + assert state.state == STATE_OFF + + +@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True) +async def test_websocket_error( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Test WebSocket connection erroring out, marking lights unavailable.""" + state = hass.states.get("light.wled_websocket") + assert state + assert state.state == STATE_ON + + connection_connected = asyncio.Future() + connection_finished = asyncio.Future() + + async def connect(callback: Callable[[WLEDDevice], None]): + connection_connected.set_result(None) + await connection_finished + + mock_wled.listen.side_effect = connect + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await connection_connected + + # Resolve Future with an error. + connection_finished.set_exception(WLEDError) + await hass.async_block_till_done() + + # Light no longer available as an error occurred + state = hass.states.get("light.wled_websocket") + assert state + assert state.state == STATE_UNAVAILABLE + + +@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True) +async def test_websocket_disconnect_on_home_assistant_stop( + hass: HomeAssistant, + init_integration: MockConfigEntry, + mock_wled: MagicMock, +) -> None: + """Ensure WebSocket is disconnected when Home Assistant stops.""" + assert mock_wled.disconnect.call_count == 1 + connection_connected = asyncio.Future() + connection_finished = asyncio.Future() + + async def connect(callback: Callable[[WLEDDevice], None]): + connection_connected.set_result(None) + await connection_finished + + mock_wled.listen.side_effect = connect + async_fire_time_changed(hass, dt_util.utcnow() + SCAN_INTERVAL) + await connection_connected + + assert mock_wled.disconnect.call_count == 1 + + hass.bus.fire(EVENT_HOMEASSISTANT_STOP) + await hass.async_block_till_done() + await hass.async_block_till_done() + assert mock_wled.disconnect.call_count == 2 diff --git a/tests/components/wled/test_init.py b/tests/components/wled/test_init.py index 50ee5520e5d..01821262389 100644 --- a/tests/components/wled/test_init.py +++ b/tests/components/wled/test_init.py @@ -1,6 +1,9 @@ """Tests for the WLED integration.""" +import asyncio +from typing import Callable from unittest.mock import AsyncMock, MagicMock, patch +import pytest from wled import WLEDConnectionError from homeassistant.components.wled.const import DOMAIN @@ -10,19 +13,36 @@ from homeassistant.core import HomeAssistant from tests.common import MockConfigEntry +@pytest.mark.parametrize("mock_wled", ["wled/rgb_websocket.json"], indirect=True) async def test_load_unload_config_entry( hass: HomeAssistant, mock_config_entry: MockConfigEntry, mock_wled: AsyncMock ) -> None: """Test the WLED configuration entry unloading.""" + connection_connected = asyncio.Future() + connection_finished = asyncio.Future() + + async def connect(callback: Callable): + connection_connected.set_result(None) + await connection_finished + + # Mock out wled.listen with a Future + mock_wled.listen.side_effect = connect + mock_config_entry.add_to_hass(hass) await hass.config_entries.async_setup(mock_config_entry.entry_id) await hass.async_block_till_done() + await connection_connected + # Ensure config entry is loaded and are connected assert mock_config_entry.state is ConfigEntryState.LOADED + assert mock_wled.connect.call_count == 1 + assert mock_wled.disconnect.call_count == 0 await hass.config_entries.async_unload(mock_config_entry.entry_id) await hass.async_block_till_done() + # Ensure everything is cleaned up nicely and are disconnected + assert mock_wled.disconnect.call_count == 1 assert not hass.data.get(DOMAIN) diff --git a/tests/fixtures/wled/rgb_websocket.json b/tests/fixtures/wled/rgb_websocket.json new file mode 100644 index 00000000000..7e37b489549 --- /dev/null +++ b/tests/fixtures/wled/rgb_websocket.json @@ -0,0 +1,289 @@ +{ + "state": { + "on": true, + "bri": 255, + "transition": 7, + "ps": -1, + "pl": -1, + "ccnf": { + "min": 1, + "max": 5, + "time": 12 + }, + "nl": { + "on": false, + "dur": 60, + "fade": true, + "mode": 1, + "tbri": 0, + "rem": -1 + }, + "udpn": { + "send": false, + "recv": true + }, + "lor": 0, + "mainseg": 0, + "seg": [ + { + "id": 0, + "start": 0, + "stop": 13, + "len": 13, + "grp": 1, + "spc": 0, + "on": true, + "bri": 255, + "col": [ + [ + 255, + 181, + 218 + ], + [ + 0, + 0, + 0 + ], + [ + 0, + 0, + 0 + ] + ], + "fx": 0, + "sx": 43, + "ix": 128, + "pal": 2, + "sel": true, + "rev": false, + "mi": false + } + ] + }, + "info": { + "ver": "0.12.0-b2", + "vid": 2103220, + "leds": { + "count": 13, + "rgbw": false, + "wv": false, + "pin": [ + 2 + ], + "pwr": 266, + "fps": 2, + "maxpwr": 1000, + "maxseg": 12, + "seglock": false + }, + "str": false, + "name": "WLED WebSocket", + "udpport": 21324, + "live": false, + "lm": "", + "lip": "", + "ws": 0, + "fxcount": 118, + "palcount": 56, + "wifi": { + "bssid": "AA:AA:AA:AA:AA:BB", + "rssi": -68, + "signal": 64, + "channel": 6 + }, + "fs": { + "u": 40, + "t": 1024, + "pmt": 1623156685 + }, + "ndc": 1, + "arch": "esp8266", + "core": "2_7_4_7", + "lwip": 1, + "freeheap": 22752, + "uptime": 258411, + "opt": 127, + "brand": "WLED", + "product": "FOSS", + "mac": "aabbccddeeff" + }, + "effects": [ + "Solid", + "Blink", + "Breathe", + "Wipe", + "Wipe Random", + "Random Colors", + "Sweep", + "Dynamic", + "Colorloop", + "Rainbow", + "Scan", + "Scan Dual", + "Fade", + "Theater", + "Theater Rainbow", + "Running", + "Saw", + "Twinkle", + "Dissolve", + "Dissolve Rnd", + "Sparkle", + "Sparkle Dark", + "Sparkle+", + "Strobe", + "Strobe Rainbow", + "Strobe Mega", + "Blink Rainbow", + "Android", + "Chase", + "Chase Random", + "Chase Rainbow", + "Chase Flash", + "Chase Flash Rnd", + "Rainbow Runner", + "Colorful", + "Traffic Light", + "Sweep Random", + "Running 2", + "Aurora", + "Stream", + "Scanner", + "Lighthouse", + "Fireworks", + "Rain", + "Tetrix", + "Fire Flicker", + "Gradient", + "Loading", + "Police", + "Police All", + "Two Dots", + "Two Areas", + "Circus", + "Halloween", + "Tri Chase", + "Tri Wipe", + "Tri Fade", + "Lightning", + "ICU", + "Multi Comet", + "Scanner Dual", + "Stream 2", + "Oscillate", + "Pride 2015", + "Juggle", + "Palette", + "Fire 2012", + "Colorwaves", + "Bpm", + "Fill Noise", + "Noise 1", + "Noise 2", + "Noise 3", + "Noise 4", + "Colortwinkles", + "Lake", + "Meteor", + "Meteor Smooth", + "Railway", + "Ripple", + "Twinklefox", + "Twinklecat", + "Halloween Eyes", + "Solid Pattern", + "Solid Pattern Tri", + "Spots", + "Spots Fade", + "Glitter", + "Candle", + "Fireworks Starburst", + "Fireworks 1D", + "Bouncing Balls", + "Sinelon", + "Sinelon Dual", + "Sinelon Rainbow", + "Popcorn", + "Drip", + "Plasma", + "Percent", + "Ripple Rainbow", + "Heartbeat", + "Pacifica", + "Candle Multi", + "Solid Glitter", + "Sunrise", + "Phased", + "Twinkleup", + "Noise Pal", + "Sine", + "Phased Noise", + "Flow", + "Chunchun", + "Dancing Shadows", + "Washing Machine", + "Candy Cane", + "Blends", + "TV Simulator", + "Dynamic Smooth" + ], + "palettes": [ + "Default", + "* Random Cycle", + "* Color 1", + "* Colors 1&2", + "* Color Gradient", + "* Colors Only", + "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", + "Aurora", + "Atlantica", + "C9 2", + "C9 New", + "Temperature", + "Aurora 2" + ] +} \ No newline at end of file