WLED WebSocket support - local push updates (#51683)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

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

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"

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

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"
 }

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):

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):

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
 
 

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

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)
 
 

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"
+  ]
+}