Migrate WLED to use DataUpdateCoordinator (#32565)

* Migrate WLED to use DataUpdateCoordinator

* Remove stale debug statement

* Process review suggestions

* Improve tests

* Improve tests

homeassistant/components/wled/__init__.py

@@ -2,35 +2,27 @@
 import asyncio
 from datetime import timedelta
 import logging
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict
 
-from wled import WLED, WLEDConnectionError, WLEDError
+from wled import WLED, Device as WLEDDevice, WLEDConnectionError, WLEDError
 
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_NAME, CONF_HOST
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.dispatcher import (
-    async_dispatcher_connect,
-    async_dispatcher_send,
-)
 from homeassistant.helpers.entity import Entity
-from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.helpers.typing import ConfigType
-from homeassistant.util import dt as dt_util
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import (
     ATTR_IDENTIFIERS,
     ATTR_MANUFACTURER,
     ATTR_MODEL,
     ATTR_SOFTWARE_VERSION,
-    DATA_WLED_CLIENT,
-    DATA_WLED_TIMER,
-    DATA_WLED_UPDATED,
     DOMAIN,
 )
 
@@ -49,22 +41,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up WLED from a config entry."""
 
     # Create WLED instance for this entry
-    session = async_get_clientsession(hass)
+    coordinator = WLEDDataUpdateCoordinator(hass, host=entry.data[CONF_HOST])
-    wled = WLED(entry.data[CONF_HOST], session=session)
+    await coordinator.async_refresh()
 
-    # Ensure we can connect and talk to it
+    if not coordinator.last_update_success:
-    try:
+        raise ConfigEntryNotReady
-        await wled.update()
-    except WLEDConnectionError as exception:
-        raise ConfigEntryNotReady from exception
 
     hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = {DATA_WLED_CLIENT: wled}
+    hass.data[DOMAIN][entry.entry_id] = coordinator
 
     # For backwards compat, set unique ID
     if entry.unique_id is None:
         hass.config_entries.async_update_entry(
-            entry, unique_id=wled.device.info.mac_address
+            entry, unique_id=coordinator.data.info.mac_address
         )
 
     # Set up all platforms for this device/entry.
@@ -73,32 +62,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             hass.config_entries.async_forward_entry_setup(entry, component)
         )
 
-    async def interval_update(now: dt_util.dt.datetime = None) -> None:
-        """Poll WLED device function, dispatches event after update."""
-        try:
-            await wled.update()
-        except WLEDError:
-            _LOGGER.debug("An error occurred while updating WLED", exc_info=True)
-
-        # Even if the update failed, we still send out the event.
-        # To allow entities to make themselves unavailable.
-        async_dispatcher_send(hass, DATA_WLED_UPDATED, entry.entry_id)
-
-    # Schedule update interval
-    hass.data[DOMAIN][entry.entry_id][DATA_WLED_TIMER] = async_track_time_interval(
-        hass, interval_update, SCAN_INTERVAL
-    )
-
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload WLED config entry."""
 
-    # Cancel update timer for this entry/device.
-    cancel_timer = hass.data[DOMAIN][entry.entry_id][DATA_WLED_TIMER]
-    cancel_timer()
-
     # Unload entities for this entry/device.
     await asyncio.gather(
         *(
@@ -115,26 +84,74 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
+def wled_exception_handler(func):
+    """Decorate WLED calls to handle WLED exceptions.
+
+    A decorator that wraps the passed in function, catches WLED errors,
+    and handles the availability of the device in the data coordinator.
+    """
+
+    async def handler(self, *args, **kwargs):
+        try:
+            await func(self, *args, **kwargs)
+            await self.coordinator.async_refresh()
+
+        except WLEDConnectionError as error:
+            _LOGGER.error("Error communicating with API: %s", error)
+            self.coordinator.last_update_success = False
+            self.coordinator.update_listeners()
+
+        except WLEDError as error:
+            _LOGGER.error("Invalid response from API: %s", error)
+
+    return handler
+
+
+class WLEDDataUpdateCoordinator(DataUpdateCoordinator):
+    """Class to manage fetching WLED data from single endpoint."""
+
+    def __init__(
+        self, hass: HomeAssistant, *, host: str,
+    ):
+        """Initialize global WLED data updater."""
+        self.wled = WLED(host, session=async_get_clientsession(hass))
+
+        super().__init__(
+            hass, _LOGGER, name=DOMAIN, update_interval=SCAN_INTERVAL,
+        )
+
+    def update_listeners(self) -> None:
+        """Call update on all listeners."""
+        for update_callback in self._listeners:
+            update_callback()
+
+    async def _async_update_data(self) -> WLEDDevice:
+        """Fetch data from WLED."""
+        try:
+            return await self.wled.update()
+        except WLEDError as error:
+            raise UpdateFailed(f"Invalid response from API: {error}")
+
+
 class WLEDEntity(Entity):
     """Defines a base WLED entity."""
 
     def __init__(
         self,
+        *,
         entry_id: str,
-        wled: WLED,
+        coordinator: WLEDDataUpdateCoordinator,
         name: str,
         icon: str,
         enabled_default: bool = True,
     ) -> None:
         """Initialize the WLED entity."""
-        self._attributes: Dict[str, Union[str, int, float]] = {}
-        self._available = True
         self._enabled_default = enabled_default
         self._entry_id = entry_id
         self._icon = icon
         self._name = name
         self._unsub_dispatcher = None
-        self.wled = wled
+        self.coordinator = coordinator
 
     @property
     def name(self) -> str:
@@ -149,7 +166,7 @@ class WLEDEntity(Entity):
     @property
     def available(self) -> bool:
         """Return True if entity is available."""
-        return self._available
+        return self.coordinator.last_update_success
 
     @property
     def entity_registry_enabled_default(self) -> bool:
@@ -161,42 +178,17 @@ class WLEDEntity(Entity):
         """Return the polling requirement of the entity."""
         return False
 
-    @property
-    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
-        """Return the state attributes of the entity."""
-        return self._attributes
-
     async def async_added_to_hass(self) -> None:
         """Connect to dispatcher listening for entity data notifications."""
-        self._unsub_dispatcher = async_dispatcher_connect(
+        self.coordinator.async_add_listener(self.async_write_ha_state)
-            self.hass, DATA_WLED_UPDATED, self._schedule_immediate_update
-        )
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect from update signal."""
-        self._unsub_dispatcher()
+        self.coordinator.async_remove_listener(self.async_write_ha_state)
-
-    @callback
-    def _schedule_immediate_update(self, entry_id: str) -> None:
-        """Schedule an immediate update of the entity."""
-        if entry_id == self._entry_id:
-            self.async_schedule_update_ha_state(True)
 
     async def async_update(self) -> None:
         """Update WLED entity."""
-        if not self.enabled:
+        await self.coordinator.async_request_refresh()
-            return
-
-        if self.wled.device is None:
-            self._available = False
-            return
-
-        self._available = True
-        await self._wled_update()
-
-    async def _wled_update(self) -> None:
-        """Update WLED entity."""
-        raise NotImplementedError()
 
 
 class WLEDDeviceEntity(WLEDEntity):
@@ -206,9 +198,9 @@ class WLEDDeviceEntity(WLEDEntity):
     def device_info(self) -> Dict[str, Any]:
         """Return device information about this WLED device."""
         return {
-            ATTR_IDENTIFIERS: {(DOMAIN, self.wled.device.info.mac_address)},
+            ATTR_IDENTIFIERS: {(DOMAIN, self.coordinator.data.info.mac_address)},
-            ATTR_NAME: self.wled.device.info.name,
+            ATTR_NAME: self.coordinator.data.info.name,
-            ATTR_MANUFACTURER: self.wled.device.info.brand,
+            ATTR_MANUFACTURER: self.coordinator.data.info.brand,
-            ATTR_MODEL: self.wled.device.info.product,
+            ATTR_MODEL: self.coordinator.data.info.product,
-            ATTR_SOFTWARE_VERSION: self.wled.device.info.version,
+            ATTR_SOFTWARE_VERSION: self.coordinator.data.info.version,
         }

homeassistant/components/wled/const.py

@@ -3,11 +3,6 @@
 # Integration domain
 DOMAIN = "wled"
 
-# Home Assistant data keys
-DATA_WLED_CLIENT = "wled_client"
-DATA_WLED_TIMER = "wled_timer"
-DATA_WLED_UPDATED = "wled_updated"
-
 # Attributes
 ATTR_COLOR_PRIMARY = "color_primary"
 ATTR_DURATION = "duration"

homeassistant/components/wled/light.py

@@ -1,8 +1,6 @@
 """Support for LED lights."""
 import logging
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
-
-from wled import WLED, Effect, WLEDError
 
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
@@ -24,7 +22,7 @@ from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import HomeAssistantType
 import homeassistant.util.color as color_util
 
-from . import WLEDDeviceEntity
+from . import WLEDDataUpdateCoordinator, WLEDDeviceEntity, wled_exception_handler
 from .const import (
     ATTR_COLOR_PRIMARY,
     ATTR_INTENSITY,
@@ -34,7 +32,6 @@ from .const import (
     ATTR_PRESET,
     ATTR_SEGMENT_ID,
     ATTR_SPEED,
-    DATA_WLED_CLIENT,
     DOMAIN,
 )
 
@@ -49,19 +46,12 @@ async def async_setup_entry(
     async_add_entities: Callable[[List[Entity], bool], None],
 ) -> None:
     """Set up WLED light based on a config entry."""
-    wled: WLED = hass.data[DOMAIN][entry.entry_id][DATA_WLED_CLIENT]
+    coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
-    # Does the WLED device support RGBW
+    lights = [
-    rgbw = wled.device.info.leds.rgbw
+        WLEDLight(entry.entry_id, coordinator, light.segment_id)
-
+        for light in coordinator.data.state.segments
-    # List of supported effects
+    ]
-    effects = wled.device.effects
-
-    # WLED supports splitting a strip in multiple segments
-    # Each segment will be a separate light in Home Assistant
-    lights = []
-    for light in wled.device.state.segments:
-        lights.append(WLEDLight(entry.entry_id, wled, light.segment_id, rgbw, effects))
 
     async_add_entities(lights, True)
 
@@ -70,50 +60,73 @@ class WLEDLight(Light, WLEDDeviceEntity):
     """Defines a WLED light."""
 
     def __init__(
-        self, entry_id: str, wled: WLED, segment: int, rgbw: bool, effects: List[Effect]
+        self, entry_id: str, coordinator: WLEDDataUpdateCoordinator, segment: int
     ):
         """Initialize WLED light."""
-        self._effects = effects
+        self._rgbw = coordinator.data.info.leds.rgbw
-        self._rgbw = rgbw
         self._segment = segment
 
-        self._brightness: Optional[int] = None
-        self._color: Optional[Tuple[float, float]] = None
-        self._effect: Optional[str] = None
-        self._state: Optional[bool] = None
-        self._white_value: Optional[int] = None
-
         # Only apply the segment ID if it is not the first segment
-        name = wled.device.info.name
+        name = coordinator.data.info.name
         if segment != 0:
             name += f" {segment}"
 
-        super().__init__(entry_id, wled, name, "mdi:led-strip-variant")
+        super().__init__(
+            entry_id=entry_id,
+            coordinator=coordinator,
+            name=name,
+            icon="mdi:led-strip-variant",
+        )
 
     @property
     def unique_id(self) -> str:
         """Return the unique ID for this sensor."""
-        return f"{self.wled.device.info.mac_address}_{self._segment}"
+        return f"{self.coordinator.data.info.mac_address}_{self._segment}"
+
+    @property
+    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
+        """Return the state attributes of the entity."""
+        playlist = self.coordinator.data.state.playlist
+        if playlist == -1:
+            playlist = None
+
+        preset = self.coordinator.data.state.preset
+        if preset == -1:
+            preset = None
+
+        return {
+            ATTR_INTENSITY: self.coordinator.data.state.segments[
+                self._segment
+            ].intensity,
+            ATTR_PALETTE: self.coordinator.data.state.segments[
+                self._segment
+            ].palette.name,
+            ATTR_PLAYLIST: playlist,
+            ATTR_PRESET: preset,
+            ATTR_SPEED: self.coordinator.data.state.segments[self._segment].speed,
+        }
 
     @property
     def hs_color(self) -> Optional[Tuple[float, float]]:
         """Return the hue and saturation color value [float, float]."""
-        return self._color
+        color = self.coordinator.data.state.segments[self._segment].color_primary
+        return color_util.color_RGB_to_hs(*color[:3])
 
     @property
     def effect(self) -> Optional[str]:
         """Return the current effect of the light."""
-        return self._effect
+        return self.coordinator.data.state.segments[self._segment].effect.name
 
     @property
     def brightness(self) -> Optional[int]:
         """Return the brightness of this light between 1..255."""
-        return self._brightness
+        return self.coordinator.data.state.brightness
 
     @property
     def white_value(self) -> Optional[int]:
         """Return the white value of this light between 0..255."""
-        return self._white_value
+        color = self.coordinator.data.state.segments[self._segment].color_primary
+        return color[-1] if self._rgbw else None
 
     @property
     def supported_features(self) -> int:
@@ -134,13 +147,14 @@ class WLEDLight(Light, WLEDDeviceEntity):
     @property
     def effect_list(self) -> List[str]:
         """Return the list of supported effects."""
-        return [effect.name for effect in self._effects]
+        return [effect.name for effect in self.coordinator.data.effects]
 
     @property
     def is_on(self) -> bool:
         """Return the state of the light."""
-        return bool(self._state)
+        return bool(self.coordinator.data.state.on)
 
+    @wled_exception_handler
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the light."""
         data = {ATTR_ON: False, ATTR_SEGMENT_ID: self._segment}
@@ -149,14 +163,9 @@ class WLEDLight(Light, WLEDDeviceEntity):
             # WLED uses 100ms per unit, so 10 = 1 second.
             data[ATTR_TRANSITION] = round(kwargs[ATTR_TRANSITION] * 10)
 
-        try:
+        await self.coordinator.wled.light(**data)
-            await self.wled.light(**data)
-            self._state = False
-        except WLEDError:
-            _LOGGER.error("An error occurred while turning off WLED light.")
-            self._available = False
-        self.async_schedule_update_ha_state()
 
+    @wled_exception_handler
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on the light."""
         data = {ATTR_ON: True, ATTR_SEGMENT_ID: self._segment}
@@ -189,8 +198,8 @@ class WLEDLight(Light, WLEDDeviceEntity):
         ):
             # WLED cannot just accept a white value, it needs the color.
             # We use the last know color in case just the white value changes.
-            if not any(x in (ATTR_COLOR_TEMP, ATTR_HS_COLOR) for x in kwargs):
+            if all(x not in (ATTR_COLOR_TEMP, ATTR_HS_COLOR) for x in kwargs):
-                hue, sat = self._color
+                hue, sat = self.hs_color
                 data[ATTR_COLOR_PRIMARY] = color_util.color_hsv_to_RGB(hue, sat, 100)
 
             # On a RGBW strip, when the color is pure white, disable the RGB LEDs in
@@ -202,56 +211,6 @@ class WLEDLight(Light, WLEDDeviceEntity):
             if ATTR_WHITE_VALUE in kwargs:
                 data[ATTR_COLOR_PRIMARY] += (kwargs[ATTR_WHITE_VALUE],)
             else:
-                data[ATTR_COLOR_PRIMARY] += (self._white_value,)
+                data[ATTR_COLOR_PRIMARY] += (self.white_value,)
 
-        try:
+        await self.coordinator.wled.light(**data)
-            await self.wled.light(**data)
-
-            self._state = True
-
-            if ATTR_BRIGHTNESS in kwargs:
-                self._brightness = kwargs[ATTR_BRIGHTNESS]
-
-            if ATTR_EFFECT in kwargs:
-                self._effect = kwargs[ATTR_EFFECT]
-
-            if ATTR_HS_COLOR in kwargs:
-                self._color = kwargs[ATTR_HS_COLOR]
-
-            if ATTR_COLOR_TEMP in kwargs:
-                self._color = color_util.color_temperature_to_hs(mireds)
-
-            if ATTR_WHITE_VALUE in kwargs:
-                self._white_value = kwargs[ATTR_WHITE_VALUE]
-
-        except WLEDError:
-            _LOGGER.error("An error occurred while turning on WLED light.")
-            self._available = False
-        self.async_schedule_update_ha_state()
-
-    async def _wled_update(self) -> None:
-        """Update WLED entity."""
-        self._brightness = self.wled.device.state.brightness
-        self._effect = self.wled.device.state.segments[self._segment].effect.name
-        self._state = self.wled.device.state.on
-
-        color = self.wled.device.state.segments[self._segment].color_primary
-        self._color = color_util.color_RGB_to_hs(*color[:3])
-        if self._rgbw:
-            self._white_value = color[-1]
-
-        playlist = self.wled.device.state.playlist
-        if playlist == -1:
-            playlist = None
-
-        preset = self.wled.device.state.preset
-        if preset == -1:
-            preset = None
-
-        self._attributes = {
-            ATTR_INTENSITY: self.wled.device.state.segments[self._segment].intensity,
-            ATTR_PALETTE: self.wled.device.state.segments[self._segment].palette.name,
-            ATTR_PLAYLIST: playlist,
-            ATTR_PRESET: preset,
-            ATTR_SPEED: self.wled.device.state.segments[self._segment].speed,
-        }

homeassistant/components/wled/sensor.py

@@ -1,7 +1,7 @@
 """Support for WLED sensors."""
 from datetime import timedelta
 import logging
-from typing import Callable, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import DATA_BYTES, DEVICE_CLASS_TIMESTAMP
@@ -9,8 +9,8 @@ from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.util.dt import utcnow
 
-from . import WLED, WLEDDeviceEntity
+from . import WLEDDataUpdateCoordinator, WLEDDeviceEntity
-from .const import ATTR_LED_COUNT, ATTR_MAX_POWER, CURRENT_MA, DATA_WLED_CLIENT, DOMAIN
+from .const import ATTR_LED_COUNT, ATTR_MAX_POWER, CURRENT_MA, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -21,12 +21,12 @@ async def async_setup_entry(
     async_add_entities: Callable[[List[Entity], bool], None],
 ) -> None:
     """Set up WLED sensor based on a config entry."""
-    wled: WLED = hass.data[DOMAIN][entry.entry_id][DATA_WLED_CLIENT]
+    coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     sensors = [
-        WLEDEstimatedCurrentSensor(entry.entry_id, wled),
+        WLEDEstimatedCurrentSensor(entry.entry_id, coordinator),
-        WLEDUptimeSensor(entry.entry_id, wled),
+        WLEDUptimeSensor(entry.entry_id, coordinator),
-        WLEDFreeHeapSensor(entry.entry_id, wled),
+        WLEDFreeHeapSensor(entry.entry_id, coordinator),
     ]
 
     async_add_entities(sensors, True)
@@ -37,30 +37,31 @@ class WLEDSensor(WLEDDeviceEntity):
 
     def __init__(
         self,
-        entry_id: str,
+        *,
-        wled: WLED,
+        coordinator: WLEDDataUpdateCoordinator,
-        name: str,
-        icon: str,
-        unit_of_measurement: str,
-        key: str,
         enabled_default: bool = True,
+        entry_id: str,
+        icon: str,
+        key: str,
+        name: str,
+        unit_of_measurement: Optional[str] = None,
     ) -> None:
         """Initialize WLED sensor."""
-        self._state = None
         self._unit_of_measurement = unit_of_measurement
         self._key = key
 
-        super().__init__(entry_id, wled, name, icon, enabled_default)
+        super().__init__(
+            entry_id=entry_id,
+            coordinator=coordinator,
+            name=name,
+            icon=icon,
+            enabled_default=enabled_default,
+        )
 
     @property
     def unique_id(self) -> str:
         """Return the unique ID for this sensor."""
-        return f"{self.wled.device.info.mac_address}_{self._key}"
+        return f"{self.coordinator.data.info.mac_address}_{self._key}"
-
-    @property
-    def state(self) -> Union[None, str, int, float]:
-        """Return the state of the sensor."""
-        return self._state
 
     @property
     def unit_of_measurement(self) -> str:
@@ -71,67 +72,73 @@ class WLEDSensor(WLEDDeviceEntity):
 class WLEDEstimatedCurrentSensor(WLEDSensor):
     """Defines a WLED estimated current sensor."""
 
-    def __init__(self, entry_id: str, wled: WLED) -> None:
+    def __init__(self, entry_id: str, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize WLED estimated current sensor."""
         super().__init__(
-            entry_id,
+            coordinator=coordinator,
-            wled,
+            entry_id=entry_id,
-            f"{wled.device.info.name} Estimated Current",
+            icon="mdi:power",
-            "mdi:power",
+            key="estimated_current",
-            CURRENT_MA,
+            name=f"{coordinator.data.info.name} Estimated Current",
-            "estimated_current",
+            unit_of_measurement=CURRENT_MA,
         )
 
-    async def _wled_update(self) -> None:
+    @property
-        """Update WLED entity."""
+    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
-        self._state = self.wled.device.info.leds.power
+        """Return the state attributes of the entity."""
-        self._attributes = {
+        return {
-            ATTR_LED_COUNT: self.wled.device.info.leds.count,
+            ATTR_LED_COUNT: self.coordinator.data.info.leds.count,
-            ATTR_MAX_POWER: self.wled.device.info.leds.max_power,
+            ATTR_MAX_POWER: self.coordinator.data.info.leds.max_power,
         }
 
+    @property
+    def state(self) -> Union[None, str, int, float]:
+        """Return the state of the sensor."""
+        return self.coordinator.data.info.leds.power
+
 
 class WLEDUptimeSensor(WLEDSensor):
     """Defines a WLED uptime sensor."""
 
-    def __init__(self, entry_id: str, wled: WLED) -> None:
+    def __init__(self, entry_id: str, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize WLED uptime sensor."""
         super().__init__(
-            entry_id,
+            coordinator=coordinator,
-            wled,
-            f"{wled.device.info.name} Uptime",
-            "mdi:clock-outline",
-            None,
-            "uptime",
             enabled_default=False,
+            entry_id=entry_id,
+            icon="mdi:clock-outline",
+            key="uptime",
+            name=f"{coordinator.data.info.name} Uptime",
         )
 
+    @property
+    def state(self) -> Union[None, str, int, float]:
+        """Return the state of the sensor."""
+        uptime = utcnow() - timedelta(seconds=self.coordinator.data.info.uptime)
+        return uptime.replace(microsecond=0).isoformat()
+
     @property
     def device_class(self) -> Optional[str]:
         """Return the class of this sensor."""
         return DEVICE_CLASS_TIMESTAMP
 
-    async def _wled_update(self) -> None:
-        """Update WLED uptime sensor."""
-        uptime = utcnow() - timedelta(seconds=self.wled.device.info.uptime)
-        self._state = uptime.replace(microsecond=0).isoformat()
-
 
 class WLEDFreeHeapSensor(WLEDSensor):
     """Defines a WLED free heap sensor."""
 
-    def __init__(self, entry_id: str, wled: WLED) -> None:
+    def __init__(self, entry_id: str, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize WLED free heap sensor."""
         super().__init__(
-            entry_id,
+            coordinator=coordinator,
-            wled,
-            f"{wled.device.info.name} Free Memory",
-            "mdi:memory",
-            DATA_BYTES,
-            "free_heap",
             enabled_default=False,
+            entry_id=entry_id,
+            icon="mdi:memory",
+            key="free_heap",
+            name=f"{coordinator.data.info.name} Free Memory",
+            unit_of_measurement=DATA_BYTES,
         )
 
-    async def _wled_update(self) -> None:
+    @property
-        """Update WLED uptime sensor."""
+    def state(self) -> Union[None, str, int, float]:
-        self._state = self.wled.device.info.free_heap
+        """Return the state of the sensor."""
+        return self.coordinator.data.info.free_heap

homeassistant/components/wled/switch.py

@@ -1,21 +1,18 @@
 """Support for WLED switches."""
 import logging
-from typing import Any, Callable, List
+from typing import Any, Callable, Dict, List, Optional
-
-from wled import WLED, WLEDError
 
 from homeassistant.components.switch import SwitchDevice
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.typing import HomeAssistantType
 
-from . import WLEDDeviceEntity
+from . import WLEDDataUpdateCoordinator, WLEDDeviceEntity, wled_exception_handler
 from .const import (
     ATTR_DURATION,
     ATTR_FADE,
     ATTR_TARGET_BRIGHTNESS,
     ATTR_UDP_PORT,
-    DATA_WLED_CLIENT,
     DOMAIN,
 )
 
@@ -30,12 +27,12 @@ async def async_setup_entry(
     async_add_entities: Callable[[List[Entity], bool], None],
 ) -> None:
     """Set up WLED switch based on a config entry."""
-    wled: WLED = hass.data[DOMAIN][entry.entry_id][DATA_WLED_CLIENT]
+    coordinator: WLEDDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     switches = [
-        WLEDNightlightSwitch(entry.entry_id, wled),
+        WLEDNightlightSwitch(entry.entry_id, coordinator),
-        WLEDSyncSendSwitch(entry.entry_id, wled),
+        WLEDSyncSendSwitch(entry.entry_id, coordinator),
-        WLEDSyncReceiveSwitch(entry.entry_id, wled),
+        WLEDSyncReceiveSwitch(entry.entry_id, coordinator),
     ]
     async_add_entities(switches, True)
 
@@ -44,132 +41,127 @@ class WLEDSwitch(WLEDDeviceEntity, SwitchDevice):
     """Defines a WLED switch."""
 
     def __init__(
-        self, entry_id: str, wled: WLED, name: str, icon: str, key: str
+        self,
+        *,
+        entry_id: str,
+        coordinator: WLEDDataUpdateCoordinator,
+        name: str,
+        icon: str,
+        key: str,
     ) -> None:
         """Initialize WLED switch."""
         self._key = key
-        self._state = False
+        super().__init__(
-        super().__init__(entry_id, wled, name, icon)
+            entry_id=entry_id, coordinator=coordinator, name=name, icon=icon
+        )
 
     @property
     def unique_id(self) -> str:
         """Return the unique ID for this sensor."""
-        return f"{self.wled.device.info.mac_address}_{self._key}"
+        return f"{self.coordinator.data.info.mac_address}_{self._key}"
-
-    @property
-    def is_on(self) -> bool:
-        """Return the state of the switch."""
-        return self._state
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn off the switch."""
-        try:
-            await self._wled_turn_off()
-            self._state = False
-        except WLEDError:
-            _LOGGER.error("An error occurred while turning off WLED switch.")
-            self._available = False
-        self.async_schedule_update_ha_state()
-
-    async def _wled_turn_off(self) -> None:
-        """Turn off the switch."""
-        raise NotImplementedError()
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn on the switch."""
-        try:
-            await self._wled_turn_on()
-            self._state = True
-        except WLEDError:
-            _LOGGER.error("An error occurred while turning on WLED switch")
-            self._available = False
-        self.async_schedule_update_ha_state()
-
-    async def _wled_turn_on(self) -> None:
-        """Turn on the switch."""
-        raise NotImplementedError()
 
 
 class WLEDNightlightSwitch(WLEDSwitch):
     """Defines a WLED nightlight switch."""
 
-    def __init__(self, entry_id: str, wled: WLED) -> None:
+    def __init__(self, entry_id: str, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize WLED nightlight switch."""
         super().__init__(
-            entry_id,
+            coordinator=coordinator,
-            wled,
+            entry_id=entry_id,
-            f"{wled.device.info.name} Nightlight",
+            icon="mdi:weather-night",
-            "mdi:weather-night",
+            key="nightlight",
-            "nightlight",
+            name=f"{coordinator.data.info.name} Nightlight",
         )
 
-    async def _wled_turn_off(self) -> None:
+    @property
-        """Turn off the WLED nightlight switch."""
+    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
-        await self.wled.nightlight(on=False)
+        """Return the state attributes of the entity."""
-
+        return {
-    async def _wled_turn_on(self) -> None:
+            ATTR_DURATION: self.coordinator.data.state.nightlight.duration,
-        """Turn on the WLED nightlight switch."""
+            ATTR_FADE: self.coordinator.data.state.nightlight.fade,
-        await self.wled.nightlight(on=True)
+            ATTR_TARGET_BRIGHTNESS: self.coordinator.data.state.nightlight.target_brightness,
-
-    async def _wled_update(self) -> None:
-        """Update WLED entity."""
-        self._state = self.wled.device.state.nightlight.on
-        self._attributes = {
-            ATTR_DURATION: self.wled.device.state.nightlight.duration,
-            ATTR_FADE: self.wled.device.state.nightlight.fade,
-            ATTR_TARGET_BRIGHTNESS: self.wled.device.state.nightlight.target_brightness,
         }
 
+    @property
+    def is_on(self) -> bool:
+        """Return the state of the switch."""
+        return bool(self.coordinator.data.state.nightlight.on)
+
+    @wled_exception_handler
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn off the WLED nightlight switch."""
+        await self.coordinator.wled.nightlight(on=False)
+
+    @wled_exception_handler
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn on the WLED nightlight switch."""
+        await self.coordinator.wled.nightlight(on=True)
+
 
 class WLEDSyncSendSwitch(WLEDSwitch):
     """Defines a WLED sync send switch."""
 
-    def __init__(self, entry_id: str, wled: WLED) -> None:
+    def __init__(self, entry_id: str, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize WLED sync send switch."""
         super().__init__(
-            entry_id,
+            coordinator=coordinator,
-            wled,
+            entry_id=entry_id,
-            f"{wled.device.info.name} Sync Send",
+            icon="mdi:upload-network-outline",
-            "mdi:upload-network-outline",
+            key="sync_send",
-            "sync_send",
+            name=f"{coordinator.data.info.name} Sync Send",
         )
 
-    async def _wled_turn_off(self) -> None:
+    @property
+    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
+        """Return the state attributes of the entity."""
+        return {ATTR_UDP_PORT: self.coordinator.data.info.udp_port}
+
+    @property
+    def is_on(self) -> bool:
+        """Return the state of the switch."""
+        return bool(self.coordinator.data.state.sync.send)
+
+    @wled_exception_handler
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the WLED sync send switch."""
-        await self.wled.sync(send=False)
+        await self.coordinator.wled.sync(send=False)
 
-    async def _wled_turn_on(self) -> None:
+    @wled_exception_handler
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on the WLED sync send switch."""
-        await self.wled.sync(send=True)
+        await self.coordinator.wled.sync(send=True)
-
-    async def _wled_update(self) -> None:
-        """Update WLED entity."""
-        self._state = self.wled.device.state.sync.send
-        self._attributes = {ATTR_UDP_PORT: self.wled.device.info.udp_port}
 
 
 class WLEDSyncReceiveSwitch(WLEDSwitch):
     """Defines a WLED sync receive switch."""
 
-    def __init__(self, entry_id: str, wled: WLED):
+    def __init__(self, entry_id: str, coordinator: WLEDDataUpdateCoordinator):
         """Initialize WLED sync receive switch."""
         super().__init__(
-            entry_id,
+            coordinator=coordinator,
-            wled,
+            entry_id=entry_id,
-            f"{wled.device.info.name} Sync Receive",
+            icon="mdi:download-network-outline",
-            "mdi:download-network-outline",
+            key="sync_receive",
-            "sync_receive",
+            name=f"{coordinator.data.info.name} Sync Receive",
         )
 
-    async def _wled_turn_off(self) -> None:
+    @property
+    def device_state_attributes(self) -> Optional[Dict[str, Any]]:
+        """Return the state attributes of the entity."""
+        return {ATTR_UDP_PORT: self.coordinator.data.info.udp_port}
+
+    @property
+    def is_on(self) -> bool:
+        """Return the state of the switch."""
+        return bool(self.coordinator.data.state.sync.receive)
+
+    @wled_exception_handler
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the WLED sync receive switch."""
-        await self.wled.sync(receive=False)
+        await self.coordinator.wled.sync(receive=False)
 
-    async def _wled_turn_on(self) -> None:
+    @wled_exception_handler
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn on the WLED sync receive switch."""
-        await self.wled.sync(receive=True)
+        await self.coordinator.wled.sync(receive=True)
-
-    async def _wled_update(self) -> None:
-        """Update WLED entity."""
-        self._state = self.wled.device.state.sync.receive
-        self._attributes = {ATTR_UDP_PORT: self.wled.device.info.udp_port}

tests/components/wled/test_init.py

@@ -1,10 +1,8 @@
 """Tests for the WLED integration."""
 import aiohttp
-from asynctest import patch
 
 from homeassistant.components.wled.const import DOMAIN
 from homeassistant.config_entries import ENTRY_STATE_SETUP_RETRY
-from homeassistant.const import STATE_ON
 from homeassistant.core import HomeAssistant
 
 from tests.components.wled import init_integration
@@ -39,30 +37,3 @@ async def test_setting_unique_id(hass, aioclient_mock):
 
     assert hass.data[DOMAIN]
     assert entry.unique_id == "aabbccddeeff"
-
-
-async def test_interval_update(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
-) -> None:
-    """Test the WLED configuration entry unloading."""
-    entry = await init_integration(hass, aioclient_mock, skip_setup=True)
-
-    interval_action = False
-
-    def async_track_time_interval(hass, action, interval):
-        nonlocal interval_action
-        interval_action = action
-
-    with patch(
-        "homeassistant.components.wled.async_track_time_interval",
-        new=async_track_time_interval,
-    ):
-        await hass.config_entries.async_setup(entry.entry_id)
-        await hass.async_block_till_done()
-
-    assert interval_action
-    await interval_action()  # pylint: disable=not-callable
-    await hass.async_block_till_done()
-
-    state = hass.states.get("light.wled_rgb_light")
-    assert state.state == STATE_ON

tests/components/wled/test_light.py

@@ -1,5 +1,6 @@
 """Tests for the WLED light platform."""
 import aiohttp
+from asynctest.mock import patch
 
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
@@ -23,7 +24,6 @@ from homeassistant.const import (
     ATTR_ICON,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
-    STATE_OFF,
     STATE_ON,
     STATE_UNAVAILABLE,
 )
@@ -79,82 +79,98 @@ async def test_rgb_light_state(
 
 
 async def test_switch_change_state(
-    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, caplog
 ) -> None:
     """Test the change of state of the WLED switches."""
     await init_integration(hass, aioclient_mock)
 
-    state = hass.states.get("light.wled_rgb_light")
+    with patch("wled.WLED.light") as light_mock:
-    assert state.state == STATE_ON
+        await hass.services.async_call(
+            LIGHT_DOMAIN,
+            SERVICE_TURN_OFF,
+            {ATTR_ENTITY_ID: "light.wled_rgb_light", ATTR_TRANSITION: 5},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        light_mock.assert_called_once_with(
+            on=False, segment_id=0, transition=50,
+        )
 
-    await hass.services.async_call(
+    with patch("wled.WLED.light") as light_mock:
-        LIGHT_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_OFF,
+            LIGHT_DOMAIN,
-        {ATTR_ENTITY_ID: "light.wled_rgb_light", ATTR_TRANSITION: 5},
+            SERVICE_TURN_ON,
-        blocking=True,
+            {
-    )
+                ATTR_BRIGHTNESS: 42,
-    await hass.async_block_till_done()
+                ATTR_EFFECT: "Chase",
-    state = hass.states.get("light.wled_rgb_light")
+                ATTR_ENTITY_ID: "light.wled_rgb_light",
-    assert state.state == STATE_OFF
+                ATTR_RGB_COLOR: [255, 0, 0],
+                ATTR_TRANSITION: 5,
+            },
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        light_mock.assert_called_once_with(
+            brightness=42,
+            color_primary=(255, 0, 0),
+            effect="Chase",
+            on=True,
+            segment_id=0,
+            transition=50,
+        )
 
-    await hass.services.async_call(
+    with patch("wled.WLED.light") as light_mock:
-        LIGHT_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_ON,
+            LIGHT_DOMAIN,
-        {
+            SERVICE_TURN_ON,
-            ATTR_BRIGHTNESS: 42,
+            {ATTR_ENTITY_ID: "light.wled_rgb_light", ATTR_COLOR_TEMP: 400},
-            ATTR_EFFECT: "Chase",
+            blocking=True,
-            ATTR_ENTITY_ID: "light.wled_rgb_light",
+        )
-            ATTR_RGB_COLOR: [255, 0, 0],
+        await hass.async_block_till_done()
-            ATTR_TRANSITION: 5,
+        light_mock.assert_called_once_with(
-        },
+            color_primary=(255, 159, 70), on=True, segment_id=0,
-        blocking=True,
+        )
-    )
-    await hass.async_block_till_done()
-
-    state = hass.states.get("light.wled_rgb_light")
-    assert state.state == STATE_ON
-    assert state.attributes.get(ATTR_BRIGHTNESS) == 42
-    assert state.attributes.get(ATTR_EFFECT) == "Chase"
-    assert state.attributes.get(ATTR_HS_COLOR) == (0.0, 100.0)
-
-    await hass.services.async_call(
-        LIGHT_DOMAIN,
-        SERVICE_TURN_ON,
-        {ATTR_ENTITY_ID: "light.wled_rgb_light", ATTR_COLOR_TEMP: 400},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    state = hass.states.get("light.wled_rgb_light")
-    assert state.state == STATE_ON
-    assert state.attributes.get(ATTR_HS_COLOR) == (28.874, 72.522)
 
 
 async def test_light_error(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, caplog
+) -> None:
+    """Test error handling of the WLED lights."""
+    aioclient_mock.post("http://example.local:80/json/state", text="", status=400)
+    await init_integration(hass, aioclient_mock)
+
+    with patch("homeassistant.components.wled.WLEDDataUpdateCoordinator.async_refresh"):
+        await hass.services.async_call(
+            LIGHT_DOMAIN,
+            SERVICE_TURN_OFF,
+            {ATTR_ENTITY_ID: "light.wled_rgb_light"},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+
+        state = hass.states.get("light.wled_rgb_light")
+        assert state.state == STATE_ON
+        assert "Invalid response from API" in caplog.text
+
+
+async def test_light_connection_error(
     hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
 ) -> None:
     """Test error handling of the WLED switches."""
     aioclient_mock.post("http://example.local:80/json/state", exc=aiohttp.ClientError)
     await init_integration(hass, aioclient_mock)
 
-    await hass.services.async_call(
+    with patch("homeassistant.components.wled.WLEDDataUpdateCoordinator.async_refresh"):
-        LIGHT_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_OFF,
+            LIGHT_DOMAIN,
-        {ATTR_ENTITY_ID: "light.wled_rgb_light"},
+            SERVICE_TURN_OFF,
-        blocking=True,
+            {ATTR_ENTITY_ID: "light.wled_rgb_light"},
-    )
+            blocking=True,
-    await hass.async_block_till_done()
+        )
-    state = hass.states.get("light.wled_rgb_light")
+        await hass.async_block_till_done()
-    assert state.state == STATE_UNAVAILABLE
 
-    await hass.services.async_call(
+        state = hass.states.get("light.wled_rgb_light")
-        LIGHT_DOMAIN,
+        assert state.state == STATE_UNAVAILABLE
-        SERVICE_TURN_ON,
-        {ATTR_ENTITY_ID: "light.wled_rgb_light_1"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    state = hass.states.get("light.wled_rgb_light_1")
-    assert state.state == STATE_UNAVAILABLE
 
 
 async def test_rgbw_light(
@@ -168,45 +184,42 @@ async def test_rgbw_light(
     assert state.attributes.get(ATTR_HS_COLOR) == (0.0, 100.0)
     assert state.attributes.get(ATTR_WHITE_VALUE) == 139
 
-    await hass.services.async_call(
+    with patch("wled.WLED.light") as light_mock:
-        LIGHT_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_ON,
+            LIGHT_DOMAIN,
-        {ATTR_ENTITY_ID: "light.wled_rgbw_light", ATTR_COLOR_TEMP: 400},
+            SERVICE_TURN_ON,
-        blocking=True,
+            {ATTR_ENTITY_ID: "light.wled_rgbw_light", ATTR_COLOR_TEMP: 400},
-    )
+            blocking=True,
-    await hass.async_block_till_done()
+        )
+        await hass.async_block_till_done()
+        light_mock.assert_called_once_with(
+            on=True, segment_id=0, color_primary=(255, 159, 70, 139),
+        )
 
-    state = hass.states.get("light.wled_rgbw_light")
+    with patch("wled.WLED.light") as light_mock:
-    assert state.state == STATE_ON
+        await hass.services.async_call(
-    assert state.attributes.get(ATTR_HS_COLOR) == (28.874, 72.522)
+            LIGHT_DOMAIN,
-    assert state.attributes.get(ATTR_WHITE_VALUE) == 139
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: "light.wled_rgbw_light", ATTR_WHITE_VALUE: 100},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        light_mock.assert_called_once_with(
+            color_primary=(255, 0, 0, 100), on=True, segment_id=0,
+        )
 
-    await hass.services.async_call(
+    with patch("wled.WLED.light") as light_mock:
-        LIGHT_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_ON,
+            LIGHT_DOMAIN,
-        {ATTR_ENTITY_ID: "light.wled_rgbw_light", ATTR_WHITE_VALUE: 100},
+            SERVICE_TURN_ON,
-        blocking=True,
+            {
-    )
+                ATTR_ENTITY_ID: "light.wled_rgbw_light",
-    await hass.async_block_till_done()
+                ATTR_RGB_COLOR: (255, 255, 255),
-
+                ATTR_WHITE_VALUE: 100,
-    state = hass.states.get("light.wled_rgbw_light")
+            },
-    assert state.state == STATE_ON
+            blocking=True,
-    assert state.attributes.get(ATTR_HS_COLOR) == (28.874, 72.522)
+        )
-    assert state.attributes.get(ATTR_WHITE_VALUE) == 100
+        await hass.async_block_till_done()
-
+        light_mock.assert_called_once_with(
-    await hass.services.async_call(
+            color_primary=(0, 0, 0, 100), on=True, segment_id=0,
-        LIGHT_DOMAIN,
+        )
-        SERVICE_TURN_ON,
-        {
-            ATTR_ENTITY_ID: "light.wled_rgbw_light",
-            ATTR_RGB_COLOR: (255, 255, 255),
-            ATTR_WHITE_VALUE: 100,
-        },
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-
-    state = hass.states.get("light.wled_rgbw_light")
-    assert state.state == STATE_ON
-    assert state.attributes.get(ATTR_HS_COLOR) == (0, 0)
-    assert state.attributes.get(ATTR_WHITE_VALUE) == 100

tests/components/wled/test_switch.py

@@ -1,5 +1,6 @@
 """Tests for the WLED switch platform."""
 import aiohttp
+from asynctest.mock import patch
 
 from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
 from homeassistant.components.wled.const import (
@@ -71,117 +72,105 @@ async def test_switch_change_state(
     await init_integration(hass, aioclient_mock)
 
     # Nightlight
-    state = hass.states.get("switch.wled_rgb_light_nightlight")
+    with patch("wled.WLED.nightlight") as nightlight_mock:
-    assert state.state == STATE_OFF
+        await hass.services.async_call(
+            SWITCH_DOMAIN,
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_nightlight"},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        nightlight_mock.assert_called_once_with(on=True)
 
-    await hass.services.async_call(
+    with patch("wled.WLED.nightlight") as nightlight_mock:
-        SWITCH_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_ON,
+            SWITCH_DOMAIN,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_nightlight"},
+            SERVICE_TURN_OFF,
-        blocking=True,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_nightlight"},
-    )
+            blocking=True,
-    await hass.async_block_till_done()
+        )
-
+        await hass.async_block_till_done()
-    state = hass.states.get("switch.wled_rgb_light_nightlight")
+        nightlight_mock.assert_called_once_with(on=False)
-    assert state.state == STATE_ON
-
-    await hass.services.async_call(
-        SWITCH_DOMAIN,
-        SERVICE_TURN_OFF,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_nightlight"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-
-    state = hass.states.get("switch.wled_rgb_light_nightlight")
-    assert state.state == STATE_OFF
 
     # Sync send
-    state = hass.states.get("switch.wled_rgb_light_sync_send")
+    with patch("wled.WLED.sync") as sync_mock:
-    assert state.state == STATE_OFF
+        await hass.services.async_call(
+            SWITCH_DOMAIN,
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_send"},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        sync_mock.assert_called_once_with(send=True)
 
-    await hass.services.async_call(
+    with patch("wled.WLED.sync") as sync_mock:
-        SWITCH_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_ON,
+            SWITCH_DOMAIN,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_send"},
+            SERVICE_TURN_OFF,
-        blocking=True,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_send"},
-    )
+            blocking=True,
-    await hass.async_block_till_done()
+        )
-
+        await hass.async_block_till_done()
-    state = hass.states.get("switch.wled_rgb_light_sync_send")
+        sync_mock.assert_called_once_with(send=False)
-    assert state.state == STATE_ON
-
-    await hass.services.async_call(
-        SWITCH_DOMAIN,
-        SERVICE_TURN_OFF,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_send"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-
-    state = hass.states.get("switch.wled_rgb_light_sync_send")
-    assert state.state == STATE_OFF
 
     # Sync receive
-    state = hass.states.get("switch.wled_rgb_light_sync_receive")
+    with patch("wled.WLED.sync") as sync_mock:
-    assert state.state == STATE_ON
+        await hass.services.async_call(
+            SWITCH_DOMAIN,
+            SERVICE_TURN_OFF,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_receive"},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+        sync_mock.assert_called_once_with(receive=False)
 
-    await hass.services.async_call(
+    with patch("wled.WLED.sync") as sync_mock:
-        SWITCH_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_OFF,
+            SWITCH_DOMAIN,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_receive"},
+            SERVICE_TURN_ON,
-        blocking=True,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_receive"},
-    )
+            blocking=True,
-    await hass.async_block_till_done()
+        )
-
+        await hass.async_block_till_done()
-    state = hass.states.get("switch.wled_rgb_light_sync_receive")
+        sync_mock.assert_called_once_with(receive=True)
-    assert state.state == STATE_OFF
-
-    await hass.services.async_call(
-        SWITCH_DOMAIN,
-        SERVICE_TURN_ON,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_receive"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-
-    state = hass.states.get("switch.wled_rgb_light_sync_receive")
-    assert state.state == STATE_ON
 
 
 async def test_switch_error(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, caplog
+) -> None:
+    """Test error handling of the WLED switches."""
+    aioclient_mock.post("http://example.local:80/json/state", text="", status=400)
+    await init_integration(hass, aioclient_mock)
+
+    with patch("homeassistant.components.wled.WLEDDataUpdateCoordinator.async_refresh"):
+        await hass.services.async_call(
+            SWITCH_DOMAIN,
+            SERVICE_TURN_ON,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_nightlight"},
+            blocking=True,
+        )
+        await hass.async_block_till_done()
+
+        state = hass.states.get("switch.wled_rgb_light_nightlight")
+        assert state.state == STATE_OFF
+        assert "Invalid response from API" in caplog.text
+
+
+async def test_switch_connection_error(
     hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
 ) -> None:
     """Test error handling of the WLED switches."""
     aioclient_mock.post("http://example.local:80/json/state", exc=aiohttp.ClientError)
     await init_integration(hass, aioclient_mock)
 
-    await hass.services.async_call(
+    with patch("homeassistant.components.wled.WLEDDataUpdateCoordinator.async_refresh"):
-        SWITCH_DOMAIN,
+        await hass.services.async_call(
-        SERVICE_TURN_ON,
+            SWITCH_DOMAIN,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_nightlight"},
+            SERVICE_TURN_ON,
-        blocking=True,
+            {ATTR_ENTITY_ID: "switch.wled_rgb_light_nightlight"},
-    )
+            blocking=True,
-    await hass.async_block_till_done()
+        )
-    state = hass.states.get("switch.wled_rgb_light_nightlight")
+        await hass.async_block_till_done()
-    assert state.state == STATE_UNAVAILABLE
 
-    await hass.services.async_call(
+        state = hass.states.get("switch.wled_rgb_light_nightlight")
-        SWITCH_DOMAIN,
+        assert state.state == STATE_UNAVAILABLE
-        SERVICE_TURN_ON,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_send"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    state = hass.states.get("switch.wled_rgb_light_sync_send")
-    assert state.state == STATE_UNAVAILABLE
-
-    await hass.services.async_call(
-        SWITCH_DOMAIN,
-        SERVICE_TURN_OFF,
-        {ATTR_ENTITY_ID: "switch.wled_rgb_light_sync_receive"},
-        blocking=True,
-    )
-    await hass.async_block_till_done()
-    state = hass.states.get("switch.wled_rgb_light_sync_receive")
-    assert state.state == STATE_UNAVAILABLE