Codex updated: Fix async/event lifecycle edge cases

tests/test_async_.py

@@ -10,6 +10,7 @@
 from zha import async_ as zha_async
 from zha.application.gateway import Gateway
 from zha.async_ import AsyncUtilMixin, ZHAJob, ZHAJobType, create_eager_task
+from zha.decorators import periodic
 
 
 @pytest.mark.parametrize("eager_start", [True, False])
@@ -483,6 +484,38 @@ async def _increment_runs_if_in_time():
     assert results == [2, 2, -1, -1]
 
 
+@pytest.mark.parametrize("limit", [0, -1])
+async def test_gather_with_limited_concurrency_rejects_non_positive_limit(
+    limit: int,
+) -> None:
+    """Test gather_with_limited_concurrency rejects a non-positive limit."""
+    task = asyncio.create_task(asyncio.sleep(0))
+    with pytest.raises(ValueError, match="limit must be > 0"):
+        await zha_async.gather_with_limited_concurrency(limit, task)
+    await task
+
+
+async def test_periodic_cancellation_propagates() -> None:
+    """Test cancellation propagates from a periodic method."""
+    started = asyncio.Event()
+    release = asyncio.Event()
+
+    class Poller:
+        @periodic((1, 1), run_immediately=True)
+        async def poll(self) -> None:
+            started.set()
+            await release.wait()
+
+    task = asyncio.create_task(Poller().poll())
+    await started.wait()
+    task.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert task.cancelled()
+
+
 async def test_create_eager_task_312(zha_gateway: Gateway) -> None:  # pylint: disable=unused-argument
     """Test create_eager_task schedules a task eagerly in the event loop.
 

tests/test_event.py

@@ -188,6 +188,97 @@ async def test_event_emit_with_context():
     async_callback.assert_awaited_once_with("test", "data")
 
 
+async def test_event_base_once_async_multiple_emits_same_tick() -> None:
+    """Test an async once listener runs once for back-to-back emits."""
+    event = EventGenerator()
+    callback = AsyncMock()
+
+    event.once("test", callback)
+    event.emit("test", "first")
+    event.emit("test", "second")
+
+    assert event._listeners == {"test": []}
+    assert len(event._event_tasks) == 1
+
+    await asyncio.gather(*event._event_tasks)
+
+    assert callback.await_args_list == [call("first")]
+    assert not event._event_tasks
+
+
+def test_event_base_once_reentrant_emit() -> None:
+    """Test a once listener runs once when an earlier listener re-emits."""
+    event = EventGenerator()
+    reentered = False
+
+    def reentrant_listener(data: str) -> None:
+        nonlocal reentered
+        if not reentered:
+            reentered = True
+            event.emit("test", "inner")
+
+    event.on_event("test", reentrant_listener)
+    callback = MagicMock()
+    event.once("test", callback)
+
+    event.emit("test", "outer")
+
+    callback.assert_called_once_with("inner")
+
+
+async def test_event_base_emit_async_callable_object() -> None:
+    """Test an async callable listener is scheduled and tracked."""
+
+    class AsyncCallable:
+        def __init__(self) -> None:
+            self.calls: list[str] = []
+
+        async def __call__(self, data: str) -> None:
+            self.calls.append(data)
+
+    event = EventGenerator()
+    callback = AsyncCallable()
+    event.on_event("test", callback)
+
+    event.emit("test", "payload")
+
+    assert len(event._event_tasks) == 1
+    await asyncio.gather(*event._event_tasks)
+
+    assert callback.calls == ["payload"]
+    assert not event._event_tasks
+
+
+@pytest.mark.parametrize("use_task", [False, True], ids=["future", "task"])
+async def test_event_base_emit_sync_listener_returning_future(
+    use_task: bool,
+) -> None:
+    """Test a Future returned by a sync listener is tracked unchanged."""
+    event = EventGenerator()
+    release = asyncio.Event()
+    if use_task:
+        future: asyncio.Future[None] = asyncio.create_task(release.wait())
+    else:
+        future = asyncio.get_running_loop().create_future()
+    callback = MagicMock(return_value=future)
+    event.on_event("test", callback)
+
+    event.emit("test", "payload")
+
+    try:
+        callback.assert_called_once_with("payload")
+        assert event._event_tasks == [future]
+    finally:
+        if use_task:
+            release.set()
+            await future
+        else:
+            future.set_result(None)
+        await asyncio.sleep(0)
+
+    assert not event._event_tasks
+
+
 def test_handle_event_protocol():
     """Test event base class."""
 

zha/async_.py

@@ -57,6 +57,9 @@ async def gather_with_limited_concurrency(
 
     From: https://stackoverflow.com/a/61478547/9127614
     """
+    if limit <= 0:
+        raise ValueError("limit must be > 0")
+
     semaphore = Semaphore(limit)
 
     async def sem_task(task: Awaitable[Any]) -> Any:

zha/decorators.py

@@ -84,7 +84,7 @@ async def wrapper(*args: Any, **kwargs: Any) -> None:
                         asyncio.current_task(),
                         method_info,
                     )
-                    break
+                    raise
                 except Exception as ex:  # pylint: disable=broad-except
                     _LOGGER.warning(
                         "[%s] Failed to poll using method %s",

zha/event.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 import dataclasses
 import inspect
 import logging
@@ -27,9 +27,15 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize event base."""
         super().__init__(*args, **kwargs)
         self._listeners: dict[str, list[EventListener]] = {}
-        self._event_tasks: list[asyncio.Task] = []
+        self._event_tasks: list[asyncio.Future[Any]] = []
         self._global_listeners: list[EventListener] = []
 
+    def _track_event_awaitable(self, awaitable: Awaitable[Any]) -> None:
+        """Schedule and track an event listener awaitable."""
+        task = asyncio.ensure_future(awaitable)
+        self._event_tasks.append(task)
+        task.add_done_callback(self._event_tasks.remove)
+
     def on_event(  # pylint: disable=invalid-name
         self, event_name: str, callback: Callable, with_context: bool = False
     ) -> Callable:
@@ -64,22 +70,17 @@ def once(
         self, event_name: str, callback: Callable, with_context: bool = False
     ) -> Callable:
         """Listen for an event exactly once."""
-        if inspect.iscoroutinefunction(callback):
-
-            async def async_event_listener(*args, **kwargs) -> None:
-                unsub()
-                task = asyncio.create_task(callback(*args, **kwargs))
-                self._event_tasks.append(task)
-                task.add_done_callback(self._event_tasks.remove)
-
-            unsub = self.on_event(
-                event_name, async_event_listener, with_context=with_context
-            )
-            return unsub
+        consumed = False
 
         def event_listener(*args, **kwargs) -> None:
+            nonlocal consumed
+            if consumed:
+                return
+            consumed = True
             unsub()
-            callback(*args, **kwargs)
+            call = callback(*args, **kwargs)
+            if inspect.isawaitable(call):
+                self._track_event_awaitable(call)
 
         unsub = self.on_event(event_name, event_listener, with_context=with_context)
         return unsub
@@ -100,10 +101,8 @@ def emit(self, event_name: str, data=None) -> None:
             else:
                 call = listener.callback(data)
 
-            if inspect.iscoroutinefunction(listener.callback):
-                task = asyncio.create_task(call)
-                self._event_tasks.append(task)
-                task.add_done_callback(self._event_tasks.remove)
+            if inspect.isawaitable(call):
+                self._track_event_awaitable(call)
 
     def _handle_event_protocol(self, event) -> None:
         """Process an event based on event protocol."""