WLED WebSocket support - local push updates (#51683)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
This commit is contained in:
Franck Nijhof 2021-06-11 20:55:08 +02:00 committed by GitHub
parent 5cc31a98e2
commit b83b82ca7d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 580 additions and 7 deletions

View File

@ -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

View File

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

View File

@ -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

View File

@ -7,5 +7,5 @@
"zeroconf": ["_wled._tcp.local."],
"codeowners": ["@frenck"],
"quality_scale": "platinum",
"iot_class": "local_polling"
"iot_class": "local_push"
}

View File

@ -40,7 +40,7 @@ async def async_setup_entry(
WLEDWifiSignalSensor(coordinator),
]
async_add_entities(sensors, True)
async_add_entities(sensors)
class WLEDEstimatedCurrentSensor(WLEDEntity, SensorEntity):

View File

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

View File

@ -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

View File

@ -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

View File

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

289
tests/fixtures/wled/rgb_websocket.json vendored Normal file
View File

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