From 31138bf7eca1e95b22f7db8a6242a4dfb3310a14 Mon Sep 17 00:00:00 2001
From: can1357 <me@can.ac>
Date: Sun, 14 Jun 2026 19:03:58 +0200
Subject: [PATCH] test(python/robomp): added cleanup guards to queue cancel and
 shutdown tests

- Wrapped queue shutdown and cancel test assertions in `try`/`finally` blocks.
- Cancelled and awaited worker and in-flight tasks in finalizers to prevent lingering background tasks.
- Handled expected `CancelledError` exceptions during task cleanup with `suppress`.
---
 python/robomp/tests/test_queue_cancel.py   | 29 +++++++++------
 python/robomp/tests/test_queue_shutdown.py | 41 +++++++++++++---------
 2 files changed, 42 insertions(+), 28 deletions(-)

diff --git a/python/robomp/tests/test_queue_cancel.py b/python/robomp/tests/test_queue_cancel.py
index ba2e402bc..d8560a6b1 100644
--- a/python/robomp/tests/test_queue_cancel.py
+++ b/python/robomp/tests/test_queue_cancel.py
@@ -9,6 +9,7 @@ real omp subprocess; that's covered by the integration smoke test.
 from __future__ import annotations
 
 import asyncio
+from contextlib import suppress
 
 import pytest
 
@@ -84,18 +85,24 @@ async def test_cancel_fires_hook_armed_by_worker(settings: Settings, db: Databas
             clear_current_event(token)
 
     worker = asyncio.create_task(fake_worker())
-    # Give the worker a tick to register.
-    for _ in range(20):
-        await asyncio.sleep(0)
-        if row.delivery_id in pool._cancel_hooks:  # noqa: SLF001 — test inspecting state
-            break
-    assert row.delivery_id in pool._cancel_hooks  # noqa: SLF001
+    try:
+        # Give the worker a tick to register.
+        for _ in range(20):
+            await asyncio.sleep(0)
+            if row.delivery_id in pool._cancel_hooks:  # noqa: SLF001 — test inspecting state
+                break
+        assert row.delivery_id in pool._cancel_hooks  # noqa: SLF001
 
-    assert await pool.cancel_event(row.delivery_id) is True
-    await asyncio.wait_for(worker, timeout=1.0)
-    assert row.delivery_id in pool._cancelled  # noqa: SLF001
-    # Hook is consumed.
-    assert row.delivery_id not in pool._cancel_hooks  # noqa: SLF001
+        assert await pool.cancel_event(row.delivery_id) is True
+        await asyncio.wait_for(worker, timeout=1.0)
+        assert row.delivery_id in pool._cancelled  # noqa: SLF001
+        # Hook is consumed.
+        assert row.delivery_id not in pool._cancel_hooks  # noqa: SLF001
+    finally:
+        if not worker.done():
+            worker.cancel()
+            with suppress(asyncio.CancelledError):
+                await worker
 
 
 @pytest.mark.asyncio
diff --git a/python/robomp/tests/test_queue_shutdown.py b/python/robomp/tests/test_queue_shutdown.py
index c67dd594b..2191a17ed 100644
--- a/python/robomp/tests/test_queue_shutdown.py
+++ b/python/robomp/tests/test_queue_shutdown.py
@@ -160,16 +160,17 @@ async def test_stop_fires_kill_hook_when_drain_exceeds_timeout(settings: Setting
     blocked = asyncio.create_task(_park())
     pool._inflight_tasks[blocked] = "d-blocked"  # noqa: SLF001
 
-    await pool.stop(drain_timeout=0.05, kill_timeout=0.05)
+    try:
+        await pool.stop(drain_timeout=0.05, kill_timeout=0.05)
 
-    assert hook_called.is_set()
-    stored = db.get_event("d-blocked")
-    assert stored is not None
-    assert stored.state == "running"
-
-    blocked.cancel()
-    with suppress(asyncio.CancelledError):
-        await blocked
+        assert hook_called.is_set()
+        stored = db.get_event("d-blocked")
+        assert stored is not None
+        assert stored.state == "running"
+    finally:
+        blocked.cancel()
+        with suppress(asyncio.CancelledError):
+            await blocked
 
 
 @pytest.mark.asyncio
@@ -230,16 +231,22 @@ async def test_stop_cancels_hookless_inflight_task(settings: Settings, db: Datab
 
     task = asyncio.create_task(stuck_pre_hook())
     pool._inflight_tasks[task] = "d-hookless"  # noqa: SLF001
-    await asyncio.wait_for(pre_hook_started.wait(), timeout=1.0)
+    try:
+        await asyncio.wait_for(pre_hook_started.wait(), timeout=1.0)
 
-    await pool.stop(drain_timeout=0.05, kill_timeout=0.2)
+        await pool.stop(drain_timeout=0.05, kill_timeout=0.2)
 
-    # Give the event loop a tick for cancellation to settle, then assert.
-    await asyncio.sleep(0)
-    assert task.done(), "stop() must terminate hookless in-flight tasks"
-    assert task.cancelled(), "hookless task must be cancelled, not left running"
-    assert reached_spawn is False, "task body must not progress past stop()"
-    assert "d-hookless" in pool._shutdown_cancelled  # noqa: SLF001
+        # Give the event loop a tick for cancellation to settle, then assert.
+        await asyncio.sleep(0)
+        assert task.done(), "stop() must terminate hookless in-flight tasks"
+        assert task.cancelled(), "hookless task must be cancelled, not left running"
+        assert reached_spawn is False, "task body must not progress past stop()"
+        assert "d-hookless" in pool._shutdown_cancelled  # noqa: SLF001
+    finally:
+        if not task.done():
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task
 
 
 @pytest.mark.asyncio