Fix tplink effect not being restored when turning back on (#68533)

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

homeassistant/components/tplink/light.py

@@ -21,7 +21,7 @@ from homeassistant.components.light import (
     LightEntity,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.color import (
     color_temperature_kelvin_to_mired as kelvin_to_mired,
@@ -43,11 +43,18 @@ async def async_setup_entry(
 ) -> None:
     """Set up switches."""
     coordinator: TPLinkDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
-    device = cast(SmartBulb, coordinator.device)
+    if coordinator.device.is_light_strip:
-    if device.is_light_strip:
+        async_add_entities(
-        async_add_entities([TPLinkSmartLightStrip(device, coordinator)])
+            [
-    elif device.is_bulb or device.is_dimmer:
+                TPLinkSmartLightStrip(
-        async_add_entities([TPLinkSmartBulb(device, coordinator)])
+                    cast(SmartLightStrip, coordinator.device), coordinator
+                )
+            ]
+        )
+    elif coordinator.device.is_bulb or coordinator.device.is_dimmer:
+        async_add_entities(
+            [TPLinkSmartBulb(cast(SmartBulb, coordinator.device), coordinator)]
+        )
 
 
 class TPLinkSmartBulb(CoordinatedTPLinkEntity, LightEntity):
@@ -72,9 +79,10 @@ class TPLinkSmartBulb(CoordinatedTPLinkEntity, LightEntity):
         else:
             self._attr_unique_id = device.mac.replace(":", "").upper()
 
-    @async_refresh_after
+    @callback
-    async def async_turn_on(self, **kwargs: Any) -> None:
+    def _async_extract_brightness_transition(
-        """Turn the light on."""
+        self, **kwargs: Any
+    ) -> tuple[int | None, int | None]:
         if (transition := kwargs.get(ATTR_TRANSITION)) is not None:
             transition = int(transition * 1_000)
 
@@ -89,35 +97,48 @@ class TPLinkSmartBulb(CoordinatedTPLinkEntity, LightEntity):
             # except when transition is defined, so we leverage that here for now.
             transition = 1
 
-        # Handle turning to temp mode
+        return brightness, transition
-        if ATTR_COLOR_TEMP in kwargs:
-            # Handle temp conversion mireds -> kelvin being slightly outside of valid range
-            kelvin = mired_to_kelvin(int(kwargs[ATTR_COLOR_TEMP]))
-            kelvin_range = self.device.valid_temperature_range
-            color_tmp = max(kelvin_range.min, min(kelvin_range.max, kelvin))
-            _LOGGER.debug("Changing color temp to %s", color_tmp)
-            await self.device.set_color_temp(
-                color_tmp, brightness=brightness, transition=transition
-            )
-            return
 
-        # Handling turning to hs color mode
+    async def _async_set_color_temp(
-        if ATTR_HS_COLOR in kwargs:
+        self, color_temp_mireds: int, brightness: int | None, transition: int | None
-            # TP-Link requires integers.
+    ) -> None:
-            hue, sat = tuple(int(val) for val in kwargs[ATTR_HS_COLOR])
+        # Handle temp conversion mireds -> kelvin being slightly outside of valid range
-            await self.device.set_hsv(hue, sat, brightness, transition=transition)
+        kelvin = mired_to_kelvin(color_temp_mireds)
-            return
+        kelvin_range = self.device.valid_temperature_range
+        color_tmp = max(kelvin_range.min, min(kelvin_range.max, kelvin))
+        _LOGGER.debug("Changing color temp to %s", color_tmp)
+        await self.device.set_color_temp(
+            color_tmp, brightness=brightness, transition=transition
+        )
 
-        if ATTR_EFFECT in kwargs:
+    async def _async_set_hsv(
-            assert isinstance(self.device, SmartLightStrip)
+        self, hs_color: tuple[int, int], brightness: int | None, transition: int | None
-            await self.device.set_effect(kwargs[ATTR_EFFECT])
+    ) -> None:
-            return
+        # TP-Link requires integers.
+        hue, sat = tuple(int(val) for val in hs_color)
+        await self.device.set_hsv(hue, sat, brightness, transition=transition)
 
+    async def _async_turn_on_with_brightness(
+        self, brightness: int | None, transition: int | None
+    ) -> None:
         # Fallback to adjusting brightness or turning the bulb on
         if brightness is not None:
             await self.device.set_brightness(brightness, transition=transition)
+            return
+        await self.device.turn_on(transition=transition)  # type: ignore[arg-type]
+
+    @async_refresh_after
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the light on."""
+        brightness, transition = self._async_extract_brightness_transition(**kwargs)
+        if ATTR_COLOR_TEMP in kwargs:
+            await self._async_set_color_temp(
+                int(kwargs[ATTR_COLOR_TEMP]), brightness, transition
+            )
+        if ATTR_HS_COLOR in kwargs:
+            await self._async_set_hsv(kwargs[ATTR_HS_COLOR], brightness, transition)
         else:
-            await self.device.turn_on(transition=transition)
+            await self._async_turn_on_with_brightness(brightness, transition)
 
     @async_refresh_after
     async def async_turn_off(self, **kwargs: Any) -> None:
@@ -191,6 +212,14 @@ class TPLinkSmartLightStrip(TPLinkSmartBulb):
 
     device: SmartLightStrip
 
+    def __init__(
+        self,
+        device: SmartLightStrip,
+        coordinator: TPLinkDataUpdateCoordinator,
+    ) -> None:
+        """Initialize the smart light strip."""
+        super().__init__(device, coordinator)
+
     @property
     def supported_features(self) -> int:
         """Flag supported features."""
@@ -209,3 +238,30 @@ class TPLinkSmartLightStrip(TPLinkSmartBulb):
         if (effect := self.device.effect) and effect["enable"]:
             return cast(str, effect["name"])
         return None
+
+    @async_refresh_after
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the light on."""
+        brightness, transition = self._async_extract_brightness_transition(**kwargs)
+        if ATTR_COLOR_TEMP in kwargs:
+            await self._async_set_color_temp(
+                int(kwargs[ATTR_COLOR_TEMP]), brightness, transition
+            )
+        elif ATTR_HS_COLOR in kwargs:
+            await self._async_set_hsv(kwargs[ATTR_HS_COLOR], brightness, transition)
+        elif ATTR_EFFECT in kwargs:
+            await self.device.set_effect(kwargs[ATTR_EFFECT])
+        elif (
+            self.device.is_off
+            and self.device.effect
+            and self.device.effect["enable"] == 0
+            and self.device.effect["name"]
+        ):
+            if not self.device.effect["custom"]:
+                await self.device.set_effect(self.device.effect["name"])
+            # The device does not remember custom effects
+            # so we must set a default value or it can never turn back on
+            else:
+                await self.device.set_hsv(0, 0, 100, transition=transition)
+        else:
+            await self._async_turn_on_with_brightness(brightness, transition)

tests/components/tplink/test_light.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from datetime import timedelta
-from unittest.mock import PropertyMock
+from unittest.mock import MagicMock, PropertyMock
 
 import pytest
 
@@ -23,7 +23,7 @@ from homeassistant.components.light import (
     DOMAIN as LIGHT_DOMAIN,
 )
 from homeassistant.components.tplink.const import DOMAIN
-from homeassistant.const import ATTR_ENTITY_ID, STATE_ON
+from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry as er
 from homeassistant.setup import async_setup_component
@@ -57,14 +57,17 @@ async def test_light_unique_id(hass: HomeAssistant) -> None:
     assert entity_registry.async_get(entity_id).unique_id == "AABBCCDDEEFF"
 
 
-@pytest.mark.parametrize("transition", [2.0, None])
+@pytest.mark.parametrize(
-async def test_color_light(hass: HomeAssistant, transition: float | None) -> None:
+    "bulb, transition", [(_mocked_bulb(), 2.0), (_mocked_smart_light_strip(), None)]
+)
+async def test_color_light(
+    hass: HomeAssistant, bulb: MagicMock, transition: float | None
+) -> None:
     """Test a color light and that all transitions are correctly passed."""
     already_migrated_config_entry = MockConfigEntry(
         domain=DOMAIN, data={}, unique_id=MAC_ADDRESS
     )
     already_migrated_config_entry.add_to_hass(hass)
-    bulb = _mocked_bulb()
     bulb.color_temp = None
     with _patch_discovery(device=bulb), _patch_single_discovery(device=bulb):
         await async_setup_component(hass, tplink.DOMAIN, {tplink.DOMAIN: {}})
@@ -194,14 +197,17 @@ async def test_color_light_no_temp(hass: HomeAssistant) -> None:
     bulb.set_hsv.reset_mock()
 
 
-@pytest.mark.parametrize("is_color", [True, False])
+@pytest.mark.parametrize(
-async def test_color_temp_light(hass: HomeAssistant, is_color: bool) -> None:
+    "bulb, is_color", [(_mocked_bulb(), True), (_mocked_smart_light_strip(), False)]
+)
+async def test_color_temp_light(
+    hass: HomeAssistant, bulb: MagicMock, is_color: bool
+) -> None:
     """Test a light."""
     already_migrated_config_entry = MockConfigEntry(
         domain=DOMAIN, data={}, unique_id=MAC_ADDRESS
     )
     already_migrated_config_entry.add_to_hass(hass)
-    bulb = _mocked_bulb()
     bulb.is_color = is_color
     bulb.color_temp = 4000
     bulb.is_variable_color_temp = True
@@ -415,7 +421,7 @@ async def test_smart_strip_effects(hass: HomeAssistant) -> None:
     strip.set_effect.assert_called_once_with("Effect2")
     strip.set_effect.reset_mock()
 
-    strip.effect = {"name": "Effect1", "enable": 0}
+    strip.effect = {"name": "Effect1", "enable": 0, "custom": 0}
     async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=10))
     await hass.async_block_till_done()
 
@@ -423,10 +429,63 @@ async def test_smart_strip_effects(hass: HomeAssistant) -> None:
     assert state.state == STATE_ON
     assert ATTR_EFFECT not in state.attributes
 
-    strip.effect_list = None
+    strip.is_off = True
+    strip.is_on = False
     async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=20))
     await hass.async_block_till_done()
 
+    state = hass.states.get(entity_id)
+    assert state.state == STATE_OFF
+    assert ATTR_EFFECT not in state.attributes
+
+    await hass.services.async_call(
+        LIGHT_DOMAIN,
+        "turn_on",
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+    strip.set_effect.assert_called_once_with("Effect1")
+    strip.set_effect.reset_mock()
+
+    strip.is_off = False
+    strip.is_on = True
+    strip.effect_list = None
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=30))
+    await hass.async_block_till_done()
+
     state = hass.states.get(entity_id)
     assert state.state == STATE_ON
     assert state.attributes[ATTR_EFFECT_LIST] is None
+
+
+async def test_smart_strip_custom_random_effect_at_start(hass: HomeAssistant) -> None:
+    """Test smart strip custom random effects at startup."""
+    already_migrated_config_entry = MockConfigEntry(
+        domain=DOMAIN, data={}, unique_id=MAC_ADDRESS
+    )
+    already_migrated_config_entry.add_to_hass(hass)
+    strip = _mocked_smart_light_strip()
+    strip.effect = {
+        "custom": 1,
+        "id": "yMwcNpLxijmoKamskHCvvravpbnIqAIN",
+        "brightness": 100,
+        "name": "Custom",
+        "enable": 0,
+    }
+    with _patch_discovery(device=strip), _patch_single_discovery(device=strip):
+        await async_setup_component(hass, tplink.DOMAIN, {tplink.DOMAIN: {}})
+        await hass.async_block_till_done()
+
+    entity_id = "light.my_bulb"
+
+    state = hass.states.get(entity_id)
+    assert state.state == STATE_ON
+    # fallback to set HSV when custom effect is not known so it does turn back on
+    await hass.services.async_call(
+        LIGHT_DOMAIN,
+        "turn_on",
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+    strip.set_hsv.assert_called_with(0, 0, 100, transition=None)
+    strip.set_hsv.reset_mock()