Make light polling and transition listener cleanup idempotent

tests/test_light.py

@@ -2186,3 +2186,125 @@ async def blocking_request(*args, **kwargs):
         await task
 
     assert entity.is_transitioning is False
+
+
+async def test_light_enable_second_pass_is_idempotent_for_polling_tasks(
+    zha_gateway: Gateway,
+) -> None:
+    """Test that enabling a light twice does not orphan the first poll task."""
+    device = await device_light_1_mock(zha_gateway)
+    entity = get_entity(device, platform=Platform.LIGHT)
+
+    # Reset to a known baseline by cancelling the task started during entity setup.
+    entity.disable()
+    await asyncio.sleep(0)
+
+    # Issue being validated:
+    # A second call to Light.enable() should be idempotent for polling lifecycle.
+    # Re-enabling an already enabled light should not leave an earlier refresh task alive.
+    #
+    # Why this is a problem:
+    # If the first refresh task is orphaned, later disable/remove only tracks the newest
+    # task reference and polling may keep running in the background unexpectedly.
+    entity.enable()
+    first_refresh_task = entity._refresh_task
+    assert first_refresh_task is not None
+
+    second_refresh_task = None
+    try:
+        entity.enable()
+        second_refresh_task = entity._refresh_task
+        assert second_refresh_task is not None
+        assert second_refresh_task is first_refresh_task
+
+        entity.disable()
+        await asyncio.sleep(0)
+
+        assert first_refresh_task.cancelled()
+        assert first_refresh_task not in entity._tracked_tasks
+    finally:
+        for task in (first_refresh_task, second_refresh_task):
+            if task is None:
+                continue
+            if not task.done():
+                task.cancel()
+            with contextlib.suppress(ValueError):
+                entity._tracked_tasks.remove(task)
+        with contextlib.suppress(asyncio.CancelledError):
+            await asyncio.gather(
+                *(
+                    t
+                    for t in (first_refresh_task, second_refresh_task)
+                    if t is not None
+                ),
+                return_exceptions=True,
+            )
+
+
+async def test_light_start_polling_replaces_completed_refresh_task(
+    zha_gateway: Gateway,
+) -> None:
+    """Test completed refresh task handles are replaced cleanly."""
+    device = await device_light_1_mock(zha_gateway)
+    entity = get_entity(device, platform=Platform.LIGHT)
+
+    # Reset baseline from entity setup.
+    entity.disable()
+    await asyncio.sleep(0)
+
+    completed_task = asyncio.create_task(asyncio.sleep(0))
+    await completed_task
+    entity._refresh_task = completed_task
+    entity._tracked_tasks.append(completed_task)
+
+    # Issue being validated:
+    # start_polling() must clean up an already-finished refresh task before creating
+    # the replacement polling task.
+    #
+    # Why this is a problem:
+    # Leaving finished task handles in tracking state causes stale bookkeeping to
+    # accumulate and can desynchronize later disable/remove cleanup behavior.
+    replacement_task = None
+    try:
+        entity.start_polling()
+        replacement_task = entity._refresh_task
+
+        assert replacement_task is not None
+        assert replacement_task is not completed_task
+        assert completed_task not in entity._tracked_tasks
+        assert replacement_task in entity._tracked_tasks
+    finally:
+        entity.disable()
+        await asyncio.sleep(0)
+        if replacement_task and replacement_task in entity._tracked_tasks:
+            entity._tracked_tasks.remove(replacement_task)
+
+
+async def test_async_unsub_transition_listener_second_pass_removes_old_handle(
+    zha_gateway: Gateway,
+) -> None:
+    """Test transition unsubscription removes the original tracked handle."""
+    device = await device_light_1_mock(zha_gateway)
+    entity = get_entity(device, platform=Platform.LIGHT)
+
+    # Issue being validated:
+    # _async_unsub_transition_listener() should remove the original listener handle from
+    # _tracked_handles when unsubscribing.
+    #
+    # Why this is a problem:
+    # Leaving cancelled timer handles in _tracked_handles leaks lifecycle bookkeeping and
+    # accumulates stale handles across repeated transition start/stop passes.
+    entity.async_transition_start_timer(transition_time=30)
+    original_listener = entity._transition_listener
+    assert original_listener is not None
+    assert original_listener in entity._tracked_handles
+
+    try:
+        entity._async_unsub_transition_listener()
+
+        assert entity._transition_listener is None
+        assert original_listener not in entity._tracked_handles
+    finally:
+        original_listener.cancel()
+        with contextlib.suppress(ValueError):
+            entity._tracked_handles.remove(original_listener)

zha/application/platforms/light/__init__.py

@@ -763,11 +763,12 @@ def async_transition_start_timer(self, transition_time) -> None:
     def _async_unsub_transition_listener(self) -> None:
         """Unsubscribe transition listener."""
         if self._transition_listener:
-            self._transition_listener.cancel()
+            transition_listener = self._transition_listener
+            transition_listener.cancel()
             self._transition_listener = None
 
             with contextlib.suppress(ValueError):
-                self._tracked_handles.remove(self._transition_listener)
+                self._tracked_handles.remove(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.
@@ -962,6 +963,13 @@ def handle_cluster_handler_set_level(self, event: LevelChangeEvent) -> None:
 
     def start_polling(self) -> None:
         """Start polling."""
+        if self._refresh_task and not self._refresh_task.done():
+            return
+
+        if self._refresh_task and self._refresh_task.done():
+            with contextlib.suppress(ValueError):
+                self._tracked_tasks.remove(self._refresh_task)
+
         self._refresh_task = self.device.gateway.async_create_background_task(
             self._refresh(),
             name=f"light_refresh_{self.unique_id}",
@@ -983,7 +991,8 @@ def disable(self) -> None:
         """Disable the entity."""
         super().disable()
         if self._refresh_task:
-            self._tracked_tasks.remove(self._refresh_task)
+            with contextlib.suppress(ValueError):
+                self._tracked_tasks.remove(self._refresh_task)
             self._refresh_task.cancel()
             self._refresh_task = None