Add switch platform to WLED integration (#28606)

* Add switch platform to WLED integration

* Use async_schedule_update_ha_state in async context

* Process review comments

homeassistant/components/wled/__init__.py

@@ -1,4 +1,5 @@
 """Support for WLED."""
+import asyncio
 from datetime import timedelta
 import logging
 from typing import Any, Dict, Optional, Union
@@ -6,6 +7,7 @@ from typing import Any, Dict, Optional, Union
 from wled import WLED, WLEDConnectionError, WLEDError
 
 from homeassistant.components.light import DOMAIN as LIGHT_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
@@ -58,9 +60,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     hass.data[DOMAIN][entry.entry_id] = {DATA_WLED_CLIENT: wled}
 
     # Set up all platforms for this device/entry.
-    hass.async_create_task(
+    for component in LIGHT_DOMAIN, SWITCH_DOMAIN:
-        hass.config_entries.async_forward_entry_setup(entry, LIGHT_DOMAIN)
+        hass.async_create_task(
-    )
+            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."""
@@ -89,7 +92,10 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     cancel_timer()
 
     # Unload entities for this entry/device.
-    await hass.config_entries.async_forward_entry_unload(entry, LIGHT_DOMAIN)
+    await asyncio.gather(
+        hass.config_entries.async_forward_entry_unload(entry, LIGHT_DOMAIN),
+        hass.config_entries.async_forward_entry_unload(entry, SWITCH_DOMAIN),
+    )
 
     # Cleanup
     del hass.data[DOMAIN][entry.entry_id]

homeassistant/components/wled/const.py

@@ -11,6 +11,7 @@ DATA_WLED_UPDATED = "wled_updated"
 # Attributes
 ATTR_COLOR_PRIMARY = "color_primary"
 ATTR_DURATION = "duration"
+ATTR_FADE = "fade"
 ATTR_IDENTIFIERS = "identifiers"
 ATTR_INTENSITY = "intensity"
 ATTR_MANUFACTURER = "manufacturer"
@@ -23,3 +24,4 @@ ATTR_SEGMENT_ID = "segment_id"
 ATTR_SOFTWARE_VERSION = "sw_version"
 ATTR_SPEED = "speed"
 ATTR_TARGET_BRIGHTNESS = "target_brightness"
+ATTR_UDP_PORT = "udp_port"

homeassistant/components/wled/switch.py

@@ -0,0 +1,175 @@
+"""Support for WLED switches."""
+import logging
+from typing import Any, Callable, List
+
+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 .const import (
+    ATTR_DURATION,
+    ATTR_FADE,
+    ATTR_TARGET_BRIGHTNESS,
+    ATTR_UDP_PORT,
+    DATA_WLED_CLIENT,
+    DOMAIN,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+PARALLEL_UPDATES = 1
+
+
+async def async_setup_entry(
+    hass: HomeAssistantType,
+    entry: ConfigEntry,
+    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]
+
+    switches = [
+        WLEDNightlightSwitch(entry.entry_id, wled),
+        WLEDSyncSendSwitch(entry.entry_id, wled),
+        WLEDSyncReceiveSwitch(entry.entry_id, wled),
+    ]
+    async_add_entities(switches, True)
+
+
+class WLEDSwitch(WLEDDeviceEntity, SwitchDevice):
+    """Defines a WLED switch."""
+
+    def __init__(
+        self, entry_id: str, wled: WLED, name: str, icon: str, key: str
+    ) -> None:
+        """Initialize WLED switch."""
+        self._key = key
+        self._state = False
+        super().__init__(entry_id, wled, name, icon)
+
+    @property
+    def unique_id(self) -> str:
+        """Return the unique ID for this sensor."""
+        return f"{self.wled.device.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:
+        """Initialize WLED nightlight switch."""
+        super().__init__(
+            entry_id,
+            wled,
+            f"{wled.device.info.name} Nightlight",
+            "mdi:weather-night",
+            "nightlight",
+        )
+
+    async def _wled_turn_off(self) -> None:
+        """Turn off the WLED nightlight switch."""
+        await self.wled.nightlight(on=False)
+
+    async def _wled_turn_on(self) -> None:
+        """Turn on the WLED nightlight switch."""
+        await self.wled.nightlight(on=True)
+
+    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,
+        }
+
+
+class WLEDSyncSendSwitch(WLEDSwitch):
+    """Defines a WLED sync send switch."""
+
+    def __init__(self, entry_id: str, wled: WLED) -> None:
+        """Initialize WLED sync send switch."""
+        super().__init__(
+            entry_id,
+            wled,
+            f"{wled.device.info.name} Sync Send",
+            "mdi:upload-network-outline",
+            "sync_send",
+        )
+
+    async def _wled_turn_off(self) -> None:
+        """Turn off the WLED sync send switch."""
+        await self.wled.sync(send=False)
+
+    async def _wled_turn_on(self) -> None:
+        """Turn on the WLED sync send switch."""
+        await self.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):
+        """Initialize WLED sync receive switch."""
+        super().__init__(
+            entry_id,
+            wled,
+            f"{wled.device.info.name} Sync Receive",
+            "mdi:download-network-outline",
+            "sync_receive",
+        )
+
+    async def _wled_turn_off(self) -> None:
+        """Turn off the WLED sync receive switch."""
+        await self.wled.sync(receive=False)
+
+    async def _wled_turn_on(self) -> None:
+        """Turn on the WLED sync receive switch."""
+        await self.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_switch.py

@@ -0,0 +1,187 @@
+"""Tests for the WLED switch platform."""
+import aiohttp
+
+from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
+from homeassistant.components.wled.const import (
+    ATTR_DURATION,
+    ATTR_FADE,
+    ATTR_TARGET_BRIGHTNESS,
+    ATTR_UDP_PORT,
+)
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_ICON,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+    STATE_OFF,
+    STATE_ON,
+    STATE_UNAVAILABLE,
+)
+from homeassistant.core import HomeAssistant
+
+from tests.components.wled import init_integration
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+
+async def test_switch_state(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the creation and values of the WLED switches."""
+    await init_integration(hass, aioclient_mock)
+
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
+
+    state = hass.states.get("switch.wled_rgb_light_nightlight")
+    assert state
+    assert state.attributes.get(ATTR_DURATION) == 60
+    assert state.attributes.get(ATTR_ICON) == "mdi:weather-night"
+    assert state.attributes.get(ATTR_TARGET_BRIGHTNESS) == 0
+    assert state.attributes.get(ATTR_FADE)
+    assert state.state == STATE_OFF
+
+    entry = entity_registry.async_get("switch.wled_rgb_light_nightlight")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_nightlight"
+
+    state = hass.states.get("switch.wled_rgb_light_sync_send")
+    assert state
+    assert state.attributes.get(ATTR_ICON) == "mdi:upload-network-outline"
+    assert state.attributes.get(ATTR_UDP_PORT) == 21324
+    assert state.state == STATE_OFF
+
+    entry = entity_registry.async_get("switch.wled_rgb_light_sync_send")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_sync_send"
+
+    state = hass.states.get("switch.wled_rgb_light_sync_receive")
+    assert state
+    assert state.attributes.get(ATTR_ICON) == "mdi:download-network-outline"
+    assert state.attributes.get(ATTR_UDP_PORT) == 21324
+    assert state.state == STATE_ON
+
+    entry = entity_registry.async_get("switch.wled_rgb_light_sync_receive")
+    assert entry
+    assert entry.unique_id == "aabbccddeeff_sync_receive"
+
+
+async def test_switch_change_state(
+    hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
+) -> None:
+    """Test the change of state of the WLED switches."""
+    await init_integration(hass, aioclient_mock)
+
+    # Nightlight
+    state = hass.states.get("switch.wled_rgb_light_nightlight")
+    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()
+
+    state = hass.states.get("switch.wled_rgb_light_nightlight")
+    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")
+    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()
+
+    state = hass.states.get("switch.wled_rgb_light_sync_send")
+    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")
+    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()
+
+    state = hass.states.get("switch.wled_rgb_light_sync_receive")
+    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
+) -> 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(
+        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_UNAVAILABLE
+
+    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()
+    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