fixes tests

Author: pcrespov

packages/common-library/src/common_library/async_tools.py

@@ -1,10 +1,13 @@
 import asyncio
 import functools
+import logging
 from collections.abc import Awaitable, Callable
 from concurrent.futures import Executor
 from inspect import isawaitable
 from typing import ParamSpec, TypeVar, overload
 
+_logger = logging.getLogger(__name__)
+
 R = TypeVar("R")
 P = ParamSpec("P")
 
@@ -67,17 +70,19 @@ async def maybe_await(
 async def cancel_and_wait(task: asyncio.Task) -> None:
     """Cancels the given task and waits for it to finish.
 
-    Accounts for the case where the parent function is being cancelled
-    and the task is cancelled as a result. In that case, it suppresses the
-    `asyncio.CancelledError` if the task was cancelled, but propagates it
-    if the task was not cancelled (i.e., it was still running when the parent
-    function was cancelled).
+    Accounts for the case where the tasks's owner function is being cancelled
     """
     task.cancel()
     try:
+        # NOTE shield ensures that cancellation of the caller function won’t stop you
+        # from observing the cancellation/finalization of task.
         await asyncio.shield(task)
     except asyncio.CancelledError:
         if not task.cancelled():
-            # parent function is being cancelled -> propagate cancellation
+            # task owner function is being cancelled -> propagate cancellation
             raise
-        # else: task was cancelled, suppress
+        # else: task cancellation is complete, we can safely ignore it
+        _logger.debug(
+            "Task %s cancellation is complete",
+            task.get_name(),
+        )

packages/common-library/tests/test_async_tools.py

@@ -13,7 +13,8 @@ def sync_function(x: int, y: int) -> int:
 
 @make_async()
 def sync_function_with_exception() -> None:
-    raise ValueError("This is an error!")
+    msg = "This is an error!"
+    raise ValueError(msg)
 
 
 @pytest.mark.asyncio
@@ -126,17 +127,23 @@ async def test_cancel_and_wait_propagates_external_cancel():
     the CancelledError is not swallowed.
     """
 
-    async def inner_coro():
+    async def coro():
         try:
-            await asyncio.sleep(10)
+            await asyncio.sleep(4)
         except asyncio.CancelledError:
-            await asyncio.sleep(0.1)  # simulate cleanup
+            await asyncio.sleep(1)  # simulate cleanup
             raise
 
-    task = asyncio.create_task(inner_coro())
+    inner_task = asyncio.create_task(coro())
 
     async def outer_coro():
-        await cancel_and_wait(task)
+        try:
+            await cancel_and_wait(inner_task)
+        except asyncio.CancelledError:
+            assert (
+                not inner_task.cancelled()
+            ), "Internal Task should not be cancelled yet (shielded)"
+            raise
 
     # Cancel the wrapper after a short delay
     outer_task = asyncio.create_task(outer_coro())
@@ -146,4 +153,13 @@ async def outer_coro():
     with pytest.raises(asyncio.CancelledError):
         await outer_task
 
-    assert task.cancelled()
+    # Ensure the task was cancelled
+    assert inner_task.cancelled() is False, "Task should not be cancelled initially"
+
+    done_event = asyncio.Event()
+
+    def on_done(_):
+        done_event.set()
+
+    inner_task.add_done_callback(on_done)
+    await done_event.wait()