TEST: Observe memory and technique worker readiness without executor capacity

doc/contributing/4_running_tests.md

@@ -114,6 +114,25 @@ timeout coverage expires the real cancellation path only after the request start
 Event-controlled analytics expiry tests freeze both deadline clocks so floating-point
 rounding cannot turn a valid query budget into a failed assertion.
 
+For thread-backed concurrency tests, signal an `asyncio.Event` from the worker with
+`loop.call_soon_threadsafe(started.set)`. Do not observe readiness through
+`asyncio.to_thread(threading_event.wait)`: that observer can queue behind the very
+worker it is waiting for when the executor is full. Use
+`wait_for_worker_ready_async(started=started, operation=task)` from
+`tests/unit/async_utils.py` when the worker must remain held during concurrent
+assertions. It builds on `wait_for_event_async` but also rejects completion after
+the readiness signal, propagating worker failure/cancellation immediately.
+
+Cancelling a task wrapping `asyncio.to_thread` does not join its thread. Prefer
+releasing the worker and joining the owned operation with `drain_tasks_async`,
+which defers caller cancellation until the owned tasks finish. Inspect task results
+separately to assert operation outcomes. If the operation can finish on cancellation
+before its thread exits, also observe actual worker completion.
+Gate timeouts should raise on the shared `TEST_WATCHDOG_SECONDS` deadlock backstop,
+not silently release workers after a short elapsed time. Exercise these patterns
+with constrained executors, including a single worker, and failure/cancellation
+paths so readiness never depends on spare executor capacity.
+
 File-surface cancellation tests hold an in-flight read at an explicit checkpoint.
 Cancellation returns before that read is released; the worker may finish that chunk,
 but must stop before starting another. Wait for worker cleanup before checking final read

tests/unit/async_utils.py

@@ -2,6 +2,8 @@
 # Licensed under the MIT license.
 
 import asyncio
+import math
+import sys
 from collections.abc import Sequence
 from contextlib import suppress
 from typing import Any, TypeVar
@@ -30,6 +32,36 @@ async def wait_for_completion_async(*, future: asyncio.Future[T], timeout: float
     return future.result()
 
 
+async def drain_tasks_async(*tasks: asyncio.Task[T], timeout: float = TEST_WATCHDOG_SECONDS) -> list[T | BaseException]:
+    """Drain released workers' owning tasks before propagating cancellation.
+
+    Callers must release worker gates first and retain tasks that still own their workers.
+    A watchdog expiry reports incomplete cleanup without cancelling those tasks.
+    Unlike ``settle_tasks_async``, cancellation must not detach a ``to_thread`` worker
+    from its owning task: cancelling that task cannot stop the underlying thread.
+    """
+    if not math.isfinite(timeout) or timeout < 0:
+        raise ValueError("The cleanup watchdog must be finite and non-negative.")
+    completion = asyncio.gather(*tasks, return_exceptions=True)
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout
+    active_exception = sys.exception()
+    cancellation = active_exception if isinstance(active_exception, asyncio.CancelledError) else None
+    while not completion.done():
+        try:
+            done, _ = await asyncio.wait({completion}, timeout=max(0, deadline - loop.time()))
+        except asyncio.CancelledError as exc:
+            if cancellation is None:
+                cancellation = exc
+            continue
+        if not done:
+            raise TimeoutError("Owned tasks did not finish before the cleanup watchdog expired.")
+    results = completion.result()
+    if cancellation is not None:
+        raise cancellation
+    return results
+
+
 async def settle_tasks_async(*tasks: asyncio.Task[Any], timeout: float = TEST_WATCHDOG_SECONDS) -> None:
     """Join owned tasks, cancelling and draining stragglers before reporting watchdog expiry.
 
@@ -72,3 +104,17 @@ async def wait_for_event_async(
         ready.cancel()
         with suppress(asyncio.CancelledError):
             await ready
+
+
+async def wait_for_worker_ready_async(
+    *, started: asyncio.Event, operation: asyncio.Future[T], timeout: float = TEST_WATCHDOG_SECONDS
+) -> None:
+    """Observe readiness while requiring the worker to remain held for concurrent assertions.
+
+    Workers signal with ``loop.call_soon_threadsafe(started.set)``. The caller still owns releasing
+    and joining the operation, including when this wait fails or is cancelled.
+    """
+    await wait_for_event_async(event=started, tasks=[operation], timeout=timeout)
+    if operation.done():
+        operation.result()
+        raise AssertionError("The operation completed before the test observed worker readiness.")

tests/unit/backend/test_technique_service.py

@@ -6,7 +6,8 @@
 import ast
 import asyncio
 import threading
-from collections.abc import Iterator
+from collections.abc import AsyncIterator, Iterator
+from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
 from typing import Any
 from unittest.mock import AsyncMock, patch
@@ -42,6 +43,14 @@
 from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse
 from pyrit.score import SubStringScorer
 from pyrit.setup.initializers.techniques import build_technique_factories
+from unit.async_utils import (
+    TEST_WATCHDOG_SECONDS,
+    drain_tasks_async,
+    get_defined_tasks,
+    wait_for_completion_async,
+    wait_for_event_async,
+    wait_for_worker_ready_async,
+)
 from unit.mocks import MockPromptTarget
 
 pytestmark = pytest.mark.usefixtures("block_dataset_provider_lookup")
@@ -59,6 +68,13 @@ def __init__(self, *, objective_target: MockPromptTarget, modes: list[_Mode]) ->
         self.mode_values = [mode.value for mode in modes]
 
 
+@pytest.fixture(params=[1, 2], ids=["one-worker", "two-workers"])
+async def worker_executor_async(request: pytest.FixtureRequest) -> AsyncIterator[ThreadPoolExecutor]:
+    with ThreadPoolExecutor(max_workers=request.param) as executor:
+        asyncio.get_running_loop().set_default_executor(executor)
+        yield executor
+
+
 @pytest.fixture
 def registry(patch_central_database: Any) -> Iterator[AttackTechniqueRegistry]:
     get_technique_service.cache_clear()
@@ -386,10 +402,12 @@ def test_rest_rejects_local_selector_collisions_and_preserves_whole_catalog(
             assert response.json() == metadata
 
 
+@pytest.mark.usefixtures("worker_executor_async")
 async def test_candidate_pool_is_not_visible_to_concurrent_catalog_reads_async(
     registry: AttackTechniqueRegistry,
 ) -> None:
-    started = threading.Event()
+    loop = asyncio.get_running_loop()
+    started = asyncio.Event()
     finish = threading.Event()
     preview = registry.create_factory(name="preview_only", attack_type="PromptSendingAttack")
     before = airt.RapidResponseTechnique
@@ -399,8 +417,8 @@ async def test_candidate_pool_is_not_visible_to_concurrent_catalog_reads_async(
     def build_technique() -> type[ScenarioTechnique]:
         factories = registry.get_factories()
         if "preview_only" in factories:
-            started.set()
-            if not finish.wait(timeout=10):
+            loop.call_soon_threadsafe(started.set)
+            if not finish.wait(timeout=TEST_WATCHDOG_SECONDS):
                 raise TimeoutError("Candidate validation did not finish")
         return registry.build_technique_class_from_factories(
             class_name="ConcurrentTechnique", factories=list(factories.values())
@@ -410,18 +428,91 @@ def build_technique() -> type[ScenarioTechnique]:
         asyncio.to_thread(registry.validate_scenario_pools, {**registry.get_factories(), preview.name: preview})
     )
     try:
-        assert await asyncio.to_thread(started.wait, 5)
+        await wait_for_worker_ready_async(started=started, operation=validation)
         assert registry.instances.get("preview_only") is None
         assert "preview_only" not in registry.get_factories()
         assert registry.catalog_revision == revision
         assert airt.RapidResponseTechnique is before
         assert "preview_only" not in {item.value for item in build_technique().get_all_techniques()}
     finally:
         finish.set()
-        await validation
+        await drain_tasks_async(validation, timeout=TEST_WATCHDOG_SECONDS)
+    validation.result()
     assert "preview_only" not in {item.value for item in build_technique().get_all_techniques()}
 
 
+@pytest.mark.usefixtures("worker_executor_async")
+async def test_candidate_validation_cleanup_defers_cancellation_until_worker_finishes_async(
+    registry: AttackTechniqueRegistry,
+) -> None:
+    loop = asyncio.get_running_loop()
+    started, finishing, torn_down = asyncio.Event(), asyncio.Event(), asyncio.Event()
+    finish, complete, finished = threading.Event(), threading.Event(), threading.Event()
+    preview = registry.create_factory(name="preview_only", attack_type="PromptSendingAttack")
+    validation: asyncio.Task[None] | None = None
+    original_builder = registry.build_technique_class_from_factories
+
+    @AttackTechniqueRegistry.cache_scenario_technique_class
+    def build_technique() -> type[ScenarioTechnique]:
+        factories = registry.get_factories()
+        if "preview_only" in factories:
+            loop.call_soon_threadsafe(started.set)
+            try:
+                assert finish.wait(timeout=TEST_WATCHDOG_SECONDS), "Validation was not released"
+                loop.call_soon_threadsafe(finishing.set)
+                assert complete.wait(timeout=TEST_WATCHDOG_SECONDS), "Validation completion was not released"
+                assert registry.build_technique_class_from_factories != original_builder
+                return registry.build_technique_class_from_factories(
+                    class_name="CleanupTechnique", factories=list(factories.values())
+                )
+            finally:
+                finished.set()
+        return original_builder(class_name="CleanupTechnique", factories=list(factories.values()))
+
+    async def owner_async() -> None:
+        nonlocal validation
+        try:
+            with patch.object(registry, "build_technique_class_from_factories", wraps=original_builder):
+                validation = asyncio.create_task(
+                    asyncio.to_thread(
+                        registry.validate_scenario_pools, {**registry.get_factories(), preview.name: preview}
+                    )
+                )
+                try:
+                    await wait_for_worker_ready_async(started=started, operation=validation)
+                finally:
+                    finish.set()
+                    await drain_tasks_async(validation, timeout=TEST_WATCHDOG_SECONDS)
+        finally:
+            torn_down.set()
+
+    owner = asyncio.create_task(owner_async())
+    try:
+        await wait_for_event_async(event=finishing, tasks=[owner])
+        assert validation is not None
+        owner.cancel("cleanup cancellation")
+        await asyncio.sleep(0)
+        owner.cancel("repeated cleanup cancellation")
+        await asyncio.sleep(0)
+        assert not owner.done()
+        assert not validation.done()
+        assert validation.cancelling() == 0
+        assert not finished.is_set()
+        assert not torn_down.is_set()
+        complete.set()
+        with pytest.raises(asyncio.CancelledError, match="^cleanup cancellation$"):
+            await wait_for_completion_async(future=owner)
+        assert validation.done()
+        assert validation.result() is None
+        assert finished.is_set()
+        assert torn_down.is_set()
+        assert registry.build_technique_class_from_factories == original_builder
+    finally:
+        finish.set()
+        complete.set()
+        await drain_tasks_async(*get_defined_tasks(owner, validation), timeout=TEST_WATCHDOG_SECONDS)
+
+
 def test_imported_selection_rebinds_only_unchanged_same_catalog_factories(
     registry: AttackTechniqueRegistry,
 ) -> None: