From da37163debdb589edd7f7ecef11d229cab01250f Mon Sep 17 00:00:00 2001
From: Roman Lutz <romanlutz13@gmail.com>
Date: Fri, 16 Oct 2026 20:10:16 -0700
Subject: [PATCH] TEST: harden media cleanup watchdog and task ownership

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: abe68236-1370-4206-b69d-2a72e999262c
---
 doc/contributing/4_running_tests.md           |  7 +++++
 tests/unit/async_utils.py                     |  5 ++-
 .../unit/backend/test_conversation_editor.py  | 31 ++++++++++++++++---
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/doc/contributing/4_running_tests.md b/doc/contributing/4_running_tests.md
index 6f47ed6afe..54baac545c 100644
--- a/doc/contributing/4_running_tests.md
+++ b/doc/contributing/4_running_tests.md
@@ -77,6 +77,13 @@ CI workers. These are deadlock backstops, not expected runtimes. Production dead
 are unchanged. Deadline and cancellation tests use explicit controls and keep their
 separate timing assertions.
 
+Async correctness tests can use `TEST_WATCHDOG_SECONDS` from `tests/unit/async_utils.py`
+for a finite 300-second deadlock backstop. Tests that start background tasks must
+release any test-controlled waits and join their tasks before removing mocks or
+disposing fixtures, including when a watchdog expires. When waiting for a task to
+signal readiness, also watch for that task finishing so an early exception is not
+hidden by the readiness timeout.
+
   * To run every test under a subdirectory:
 
       ```bash
diff --git a/tests/unit/async_utils.py b/tests/unit/async_utils.py
index 77f47db423..92e16d7f12 100644
--- a/tests/unit/async_utils.py
+++ b/tests/unit/async_utils.py
@@ -6,13 +6,16 @@
 
 T = TypeVar("T")
 
+# Correctness watchdogs allow for stalled CI workers without disabling deadlock detection.
+TEST_WATCHDOG_SECONDS = 300.0
+
 
 def get_defined_tasks(*tasks: asyncio.Task[T] | None) -> list[asyncio.Task[T]]:
     """Retain owned tasks even when a test fails before creating all of them."""
     return [task for task in tasks if task is not None]
 
 
-async def wait_for_completion_async(*, future: asyncio.Future[T], timeout: float = 300) -> T:
+async def wait_for_completion_async(*, future: asyncio.Future[T], timeout: float = TEST_WATCHDOG_SECONDS) -> T:
     """Bound a test wait without injecting cancellation into the operation under test.
 
     The default is a deadlock backstop rather than a latency assertion, so it stays finite to keep a
diff --git a/tests/unit/backend/test_conversation_editor.py b/tests/unit/backend/test_conversation_editor.py
index 14db53dacc..be702d3780 100644
--- a/tests/unit/backend/test_conversation_editor.py
+++ b/tests/unit/backend/test_conversation_editor.py
@@ -46,6 +46,7 @@
 from pyrit.models.target.target_capabilities import TargetCapabilities
 from pyrit.prompt_target import OpenAIResponseTarget
 from pyrit.score.scorer_prompt_validator import ScorerPromptValidator
+from unit.async_utils import TEST_WATCHDOG_SECONDS, wait_for_completion_async
 from unit.backend.mocks import _settle_send_async, message_send_lifecycle_async
 from unit.mocks import MockPromptTarget, get_mock_target_identifier, run_memory_session_async
 
@@ -1305,11 +1306,31 @@ async def interrupted_write_async(*args: Any, **kwargs: Any) -> None:
 
         with patch.object(storage, "write_file_async", side_effect=interrupted_write_async):
             task = asyncio.create_task(AttackService().save_conversation_async(request=request))
-            await asyncio.wait_for(started.wait(), timeout=5)
-            if cancel:
-                task.cancel()
+            started_task = asyncio.create_task(started.wait())
+            try:
+                done, _ = await asyncio.wait(
+                    {task, started_task},
+                    timeout=TEST_WATCHDOG_SECONDS,
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                if not done:
+                    raise TimeoutError("The media write did not start before the test watchdog expired.")
+                if not started.is_set():
+                    await task
+                    pytest.fail("The conversation save finished without writing media.")
+                if cancel:
+                    task.cancel()
+                    finish.set()
+                with pytest.raises(asyncio.CancelledError if cancel else OSError):
+                    await wait_for_completion_async(future=task)
+            finally:
+                if not task.done():
+                    task.cancel()
+                # The shielded writer must finish before the save task can clean up.
                 finish.set()
-            with pytest.raises(asyncio.CancelledError if cancel else OSError):
-                await task
+                started_task.cancel()
+                await wait_for_completion_async(
+                    future=asyncio.gather(task, started_task, return_exceptions=True),
+                )
         assert not list(Path(sqlite_instance.results_path).rglob("*.txt"))
         assert (await sqlite_instance.get_attack_results_async()) == []