Fix stuck light transitioning flag after task cancellation

tests/test_light.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
 from typing import Any
 from unittest.mock import AsyncMock, call, patch, sentinel
@@ -1976,3 +1977,81 @@ async def test_light_state_restoration(zha_gateway: Gateway) -> None:
     assert entity.state["xy_color"] == (1, 2)
     assert entity.state["color_mode"] == ColorMode.XY
     assert entity.state["effect"] == "colorloop"
+
+
+async def test_turn_on_cancellation_cleans_up_transition_flag(
+    zha_gateway: Gateway,
+) -> None:
+    """Test that task cancellation resets the transitioning flag.
+
+    When a mode:restart automation cancels the task, the light must not be
+    stuck ignoring attribute reports indefinitely.
+    """
+    device = await device_light_1_mock(zha_gateway)
+    entity = get_entity(device, platform=Platform.LIGHT)
+
+    cluster_level = device.device.endpoints[1].level
+
+    # Make the level cluster block indefinitely so we can cancel the task while
+    # it is suspended at the first await inside async_turn_on.
+    blocked: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+    original_request = cluster_level.request
+
+    async def blocking_request(*args, **kwargs):
+        await blocked
+        return await original_request(*args, **kwargs)
+
+    cluster_level.request = AsyncMock(side_effect=blocking_request)
+
+    # Start turn_on as a separate task, mirroring what HA does for automations.
+    task = asyncio.ensure_future(entity.async_turn_on(brightness=200, transition=1))
+
+    # Yield control so the task can run until it suspends on the cluster call.
+    await asyncio.sleep(0)
+    await asyncio.sleep(0)
+
+    # The transitioning flag must be set now, with no timer running yet.
+    assert entity.is_transitioning is True
+    assert entity._transition_listener is None
+
+    # Cancel the task (what mode:restart does to the running automation task).
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+
+    # The finally block must have cleared the flag.
+    assert entity.is_transitioning is False
+
+
+async def test_turn_off_cancellation_cleans_up_transition_flag(
+    zha_gateway: Gateway,
+) -> None:
+    """Test that task cancellation during async_turn_off resets the transitioning flag."""
+    device = await device_light_1_mock(zha_gateway)
+    entity = get_entity(device, platform=Platform.LIGHT)
+
+    cluster_on_off = device.device.endpoints[1].on_off
+
+    # Make the on/off cluster block indefinitely so we can cancel mid-turn-off.
+    blocked: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+    original_request = cluster_on_off.request
+
+    async def blocking_request(*args, **kwargs):
+        await blocked
+        return await original_request(*args, **kwargs)
+
+    cluster_on_off.request = AsyncMock(side_effect=blocking_request)
+
+    task = asyncio.ensure_future(entity.async_turn_off())
+
+    await asyncio.sleep(0)
+    await asyncio.sleep(0)
+
+    assert entity.is_transitioning is True
+    assert entity._transition_listener is None
+
+    task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await task
+
+    assert entity.is_transitioning is False

zha/application/platforms/light/__init__.py

@@ -372,6 +372,42 @@ async def async_turn_on(
         if set_transition_flag:
             self.async_transition_set_flag()
 
+        try:
+            await self._async_turn_on_impl(
+                transition=transition,
+                brightness=brightness,
+                effect=effect,
+                flash=flash,
+                color_temp=color_temp,
+                xy_color=xy_color,
+                duration=duration,
+                execute_if_off_supported=execute_if_off_supported,
+                brightness_supported=brightness_supported,
+                set_transition_flag=set_transition_flag,
+                transition_time=transition_time,
+            )
+        finally:
+            # If the task was cancelled (e.g. by a mode: restart automation) before
+            # the transition timer was started, clean up the transitioning flag so
+            # the light does not get stuck in a transitioning state indefinitely.
+            self._async_cleanup_transition_if_stuck(set_transition_flag)
+
+    async def _async_turn_on_impl(  # noqa: C901
+        self,
+        *,
+        transition: float | None,
+        brightness: int | None,
+        effect: str | None,
+        flash: FlashMode | None,
+        color_temp: int | None,
+        xy_color: tuple[int, int] | None,
+        duration: float,
+        execute_if_off_supported: bool,
+        brightness_supported: bool,
+        set_transition_flag: bool,
+        transition_time: float,
+    ) -> None:
+        """Implement the turn on logic."""
         # If the light is currently off but a turn_on call with a color/temperature is
         # sent, the light needs to be turned on first at a low brightness level where
         # the light is immediately transitioned to the correct color. Afterwards, the
@@ -527,6 +563,11 @@ async def async_turn_on(
             )
             t_log["move_to_level_if_color"] = result
             if result[1] is not Status.SUCCESS:
+                # Second 'move to level' call failed; the light is on but at the
+                # wrong brightness. If no previous timer is running, unset the flag
+                # immediately so attribute reports are not ignored indefinitely.
+                if set_transition_flag and not self._transition_listener:
+                    self.async_transition_complete()
                 self.debug("turned on: %s", t_log)
                 return
             self._state = bool(level)
@@ -594,41 +635,48 @@ async def async_turn_off(self, *, transition: float | None = None) -> None:
         if self._zha_config_enable_light_transitioning_flag:
             self.async_transition_set_flag()
 
-        # is not none looks odd here, but it will override built in bulb
-        # transition times if we pass 0 in here
-        if transition is not None and brightness_supported:
-            assert self._level_cluster_handler is not None
+        try:
+            # is not none looks odd here, but it will override built in bulb
+            # transition times if we pass 0 in here
+            if transition is not None and brightness_supported:
+                assert self._level_cluster_handler is not None
 
-            result = await self._level_cluster_handler.move_to_level_with_on_off(
-                level=0,
-                transition_time=int(
-                    10 * (transition or self._DEFAULT_MIN_TRANSITION_TIME)
-                ),
-            )
-        else:
-            assert self._on_off_cluster_handler is not None
-            result = await self._on_off_cluster_handler.off()
-
-        # Pause parsing attribute reports until transition is complete
-        if self._zha_config_enable_light_transitioning_flag:
-            self.async_transition_start_timer(transition_time)
-        self.debug("turned off: %s", result)
-        if result[1] is not Status.SUCCESS:
-            return
-        self._state = False
-
-        if brightness_supported and not self._off_with_transition:
-            # store current brightness so that the next turn_on uses it:
-            # when using "enhanced turn on"
-            self._off_brightness = self._brightness
-            if transition is not None:
-                # save for when calling turn_on without a brightness:
-                # current_level is set to 1 after transitioning to level 0,
-                # needed for correct state with light groups
-                self._brightness = 1
-                self._off_with_transition = transition is not None
+                result = await self._level_cluster_handler.move_to_level_with_on_off(
+                    level=0,
+                    transition_time=int(
+                        10 * (transition or self._DEFAULT_MIN_TRANSITION_TIME)
+                    ),
+                )
+            else:
+                assert self._on_off_cluster_handler is not None
+                result = await self._on_off_cluster_handler.off()
 
-        self.maybe_emit_state_changed_event()
+            # Pause parsing attribute reports until transition is complete
+            if self._zha_config_enable_light_transitioning_flag:
+                self.async_transition_start_timer(transition_time)
+            self.debug("turned off: %s", result)
+            if result[1] is not Status.SUCCESS:
+                return
+            self._state = False
+
+            if brightness_supported and not self._off_with_transition:
+                # store current brightness so that the next turn_on uses it:
+                # when using "enhanced turn on"
+                self._off_brightness = self._brightness
+                if transition is not None:
+                    # save for when calling turn_on without a brightness:
+                    # current_level is set to 1 after transitioning to level 0,
+                    # needed for correct state with light groups
+                    self._brightness = 1
+                    self._off_with_transition = transition is not None
+
+            self.maybe_emit_state_changed_event()
+        finally:
+            # If the task was cancelled before the transition timer was started,
+            # clean up the transitioning flag so the light does not get stuck.
+            self._async_cleanup_transition_if_stuck(
+                self._zha_config_enable_light_transitioning_flag
+            )
 
     async def async_handle_color_commands(
         self,
@@ -719,6 +767,14 @@ def _async_unsub_transition_listener(self) -> None:
             with contextlib.suppress(ValueError):
                 self._tracked_handles.remove(self._transition_listener)
 
+    def _async_cleanup_transition_if_stuck(self, guarded: bool) -> None:
+        """Call async_transition_complete if the flag is set but no timer is running.
+
+        Used in finally blocks to handle task cancellation gracefully.
+        """
+        if guarded and self._transitioning_individual and not self._transition_listener:
+            self.async_transition_complete()
+
     def async_transition_complete(self, _=None) -> None:
         """Set _transitioning_individual to False and write HA state."""
         self.debug("transition complete - future attribute reports will write HA state")