From aaf05cdf56a11f7151d74cb39e3f6d74d6d17da9 Mon Sep 17 00:00:00 2001 From: David Mulcahey Date: Tue, 14 Jul 2026 09:26:58 -0400 Subject: [PATCH] Fix async and event lifecycle edge cases --- tests/test_async_.py | 33 ++++++++++++++++ tests/test_event.py | 91 ++++++++++++++++++++++++++++++++++++++++++++ zha/async_.py | 3 ++ zha/decorators.py | 2 +- zha/event.py | 37 +++++++++--------- 5 files changed, 146 insertions(+), 20 deletions(-) diff --git a/tests/test_async_.py b/tests/test_async_.py index 9407e6fc3..c8709ccc3 100644 --- a/tests/test_async_.py +++ b/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. diff --git a/tests/test_event.py b/tests/test_event.py index 7c27543a7..9ee8e2b79 100644 --- a/tests/test_event.py +++ b/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.""" diff --git a/zha/async_.py b/zha/async_.py index 99a72bc1a..f6ad19576 100644 --- a/zha/async_.py +++ b/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: diff --git a/zha/decorators.py b/zha/decorators.py index 696ea80d1..7fb64df1e 100644 --- a/zha/decorators.py +++ b/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", diff --git a/zha/event.py b/zha/event.py index 6a31f775b..9ec51d294 100644 --- a/zha/event.py +++ b/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."""