Add support for changing Magic Home socket power restore state (#62301)

homeassistant/components/flux_led/__init__.py

@@ -45,7 +45,7 @@ PLATFORMS_BY_TYPE: Final = {
         Platform.NUMBER,
         Platform.SWITCH,
     ],
-    DeviceType.Switch: [Platform.BUTTON, Platform.SWITCH],
+    DeviceType.Switch: [Platform.BUTTON, Platform.SELECT, Platform.SWITCH],
 }
 DISCOVERY_INTERVAL: Final = timedelta(minutes=15)
 REQUEST_REFRESH_DELAY: Final = 1.5

homeassistant/components/flux_led/select.py

@@ -0,0 +1,67 @@
+"""Support for Magic Home select."""
+from __future__ import annotations
+
+from flux_led.aio import AIOWifiLedBulb
+from flux_led.protocol import PowerRestoreState
+
+from homeassistant import config_entries
+from homeassistant.components.select import SelectEntity
+from homeassistant.const import CONF_NAME
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.entity import EntityCategory
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import DOMAIN
+from .coordinator import FluxLedUpdateCoordinator
+from .entity import FluxBaseEntity
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: config_entries.ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the Flux selects."""
+    coordinator: FluxLedUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities([FluxPowerState(coordinator.device, entry)])
+
+
+def _human_readable_option(const_option: str) -> str:
+    return const_option.replace("_", " ").title()
+
+
+class FluxPowerState(FluxBaseEntity, SelectEntity):
+    """Representation of a Flux power restore state option."""
+
+    _attr_should_poll = False
+
+    def __init__(
+        self,
+        device: AIOWifiLedBulb,
+        entry: config_entries.ConfigEntry,
+    ) -> None:
+        """Initialize the power state select."""
+        super().__init__(device, entry)
+        self._attr_entity_category = EntityCategory.CONFIG
+        self._attr_name = f"{entry.data[CONF_NAME]} Power Restored"
+        if entry.unique_id:
+            self._attr_unique_id = f"{entry.unique_id}_power_restored"
+        self._name_to_state = {
+            _human_readable_option(option.name): option for option in PowerRestoreState
+        }
+        self._attr_options = list(self._name_to_state)
+        self._async_set_current_option_from_device()
+
+    @callback
+    def _async_set_current_option_from_device(self) -> None:
+        """Set the option from the current power state."""
+        restore_states = self._device.power_restore_states
+        assert restore_states is not None
+        assert restore_states.channel1 is not None
+        self._attr_current_option = _human_readable_option(restore_states.channel1.name)
+
+    async def async_select_option(self, option: str) -> None:
+        """Change the power state."""
+        await self._device.async_set_power_restore(channel1=self._name_to_state[option])
+        self._async_set_current_option_from_device()
+        self.async_write_ha_state()

tests/components/flux_led/__init__.py

@@ -14,7 +14,7 @@ from flux_led.const import (
     COLOR_MODE_RGB as FLUX_COLOR_MODE_RGB,
 )
 from flux_led.models_db import MODEL_MAP
-from flux_led.protocol import LEDENETRawState
+from flux_led.protocol import LEDENETRawState, PowerRestoreState, PowerRestoreStates
 from flux_led.scanner import FluxLEDDiscovery
 
 from homeassistant.components import dhcp
@@ -124,9 +124,16 @@ def _mocked_switch() -> AIOWifiLedBulb:
         switch.data_receive_callback = callback
 
     switch.device_type = DeviceType.Switch
+    switch.power_restore_states = PowerRestoreStates(
+        channel1=PowerRestoreState.LAST_STATE,
+        channel2=PowerRestoreState.LAST_STATE,
+        channel3=PowerRestoreState.LAST_STATE,
+        channel4=PowerRestoreState.LAST_STATE,
+    )
     switch.requires_turn_on = True
     switch.async_reboot = AsyncMock()
     switch.async_setup = AsyncMock(side_effect=_save_setup_callback)
+    switch.async_set_power_restore = AsyncMock()
     switch.async_stop = AsyncMock()
     switch.async_update = AsyncMock()
     switch.async_turn_off = AsyncMock()

tests/components/flux_led/test_select.py

@@ -0,0 +1,49 @@
+"""Tests for select platform."""
+from flux_led.protocol import PowerRestoreState
+
+from homeassistant.components import flux_led
+from homeassistant.components.flux_led.const import DOMAIN
+from homeassistant.components.select import DOMAIN as SELECT_DOMAIN
+from homeassistant.const import ATTR_ENTITY_ID, ATTR_OPTION, CONF_HOST, CONF_NAME
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from . import (
+    DEFAULT_ENTRY_TITLE,
+    IP_ADDRESS,
+    MAC_ADDRESS,
+    _mocked_switch,
+    _patch_discovery,
+    _patch_wifibulb,
+)
+
+from tests.common import MockConfigEntry
+
+
+async def test_switch_power_restore_state(hass: HomeAssistant) -> None:
+    """Test a smart plug power restore state."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={CONF_HOST: IP_ADDRESS, CONF_NAME: DEFAULT_ENTRY_TITLE},
+        unique_id=MAC_ADDRESS,
+    )
+    config_entry.add_to_hass(hass)
+    switch = _mocked_switch()
+    with _patch_discovery(), _patch_wifibulb(device=switch):
+        await async_setup_component(hass, flux_led.DOMAIN, {flux_led.DOMAIN: {}})
+        await hass.async_block_till_done()
+
+    entity_id = "select.bulb_rgbcw_ddeeff_power_restored"
+
+    state = hass.states.get(entity_id)
+    assert state.state == "Last State"
+
+    await hass.services.async_call(
+        SELECT_DOMAIN,
+        "select_option",
+        {ATTR_ENTITY_ID: entity_id, ATTR_OPTION: "Always On"},
+        blocking=True,
+    )
+    switch.async_set_power_restore.assert_called_once_with(
+        channel1=PowerRestoreState.ALWAYS_ON
+    )