diff --git a/doc/contributing/4_running_tests.md b/doc/contributing/4_running_tests.md index fdfefc7133..d109e2be53 100644 --- a/doc/contributing/4_running_tests.md +++ b/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 diff --git a/tests/unit/async_utils.py b/tests/unit/async_utils.py index f7135ecef4..b2b4d1010f 100644 --- a/tests/unit/async_utils.py +++ b/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.") diff --git a/tests/unit/backend/test_technique_service.py b/tests/unit/backend/test_technique_service.py index 255c3a476e..b89b262dcd 100644 --- a/tests/unit/backend/test_technique_service.py +++ b/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,7 +428,7 @@ 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 @@ -418,10 +436,83 @@ def build_technique() -> type[ScenarioTechnique]: 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: diff --git a/tests/unit/memory/test_async_memory.py b/tests/unit/memory/test_async_memory.py index 77321af323..6fce10b116 100644 --- a/tests/unit/memory/test_async_memory.py +++ b/tests/unit/memory/test_async_memory.py @@ -5,7 +5,8 @@ import asyncio import inspect import threading -from collections.abc import Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Sequence +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -20,6 +21,19 @@ from pyrit.common.async_compatibility import run_legacy_sync_async from pyrit.memory import AzureSQLMemory, MemoryInterface, SQLiteMemory from pyrit.models import Conversation, Message, MessagePiece, SeedPrompt +from unit.async_utils import ( + TEST_WATCHDOG_SECONDS, + drain_tasks_async, + get_defined_tasks, + wait_for_worker_ready_async, +) + + +@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 @@ -187,14 +201,16 @@ async def test_failed_first_initialization_can_retry() -> None: await memory.dispose_engine_async() +@pytest.mark.usefixtures("worker_executor_async") async def test_concurrent_initialization_is_rejected_without_disposing_first_attempt() -> None: memory = SQLiteMemory.__new__(SQLiteMemory) memory.__init__(db_path=":memory:", _defer_initialization=True) - started, release = threading.Event(), threading.Event() + loop = asyncio.get_running_loop() + started, release = asyncio.Event(), threading.Event() def initialize_schema() -> None: - started.set() - if not release.wait(timeout=5): + loop.call_soon_threadsafe(started.set) + if not release.wait(timeout=TEST_WATCHDOG_SECONDS): raise RuntimeError("Initialization was not released") try: @@ -204,13 +220,14 @@ def initialize_schema() -> None: ): task = asyncio.create_task(memory.initialize_async()) try: - assert await asyncio.to_thread(started.wait, 5) + await wait_for_worker_ready_async(started=started, operation=task) with pytest.raises(RuntimeError, match="already in progress"): await memory.initialize_async() dispose.assert_not_awaited() finally: release.set() - await task + await drain_tasks_async(task) + task.result() finally: await memory.dispose_engine_async() @@ -238,13 +255,15 @@ async def test_async_seed_operations(sqlite_instance: SQLiteMemory) -> None: assert (await sqlite_instance.get_seeds_async(dataset_name="async"))[0].value == "replacement" +@pytest.mark.usefixtures("worker_executor_async") async def test_async_database_wait_does_not_block_loop(sqlite_instance: SQLiteMemory) -> None: - started = threading.Event() + loop = asyncio.get_running_loop() + started = asyncio.Event() release = threading.Event() def wait_in_database() -> int: - started.set() - if not release.wait(timeout=10): + loop.call_soon_threadsafe(started.set) + if not release.wait(timeout=TEST_WATCHDOG_SECONDS): raise RuntimeError("Database wait was not released") return 1 @@ -257,12 +276,14 @@ def wait_in_database() -> int: ) task = asyncio.create_task(session.execute(text("SELECT wait_in_database()"))) try: - assert await asyncio.to_thread(started.wait, 5) + await wait_for_worker_ready_async(started=started, operation=task) assert not task.done() await asyncio.sleep(0) finally: release.set() - assert (await task).scalar_one() == 1 + await drain_tasks_async(task) + result = task.result() + assert result.scalar_one() == 1 async def test_azure_auth_wait_is_async_and_credential_is_closed() -> None: @@ -286,7 +307,8 @@ async def get_token_async(scope: str) -> AccessToken: assert not task.done() finally: release.set() - assert await task is token + await drain_tasks_async(task) + assert task.result() is token await auth.close_async() credential.close.assert_awaited_once() @@ -364,49 +386,55 @@ async def test_azure_closed_loop_drops_both_owned_resources() -> None: @pytest.mark.parametrize("cancellations", [1, 2]) +@pytest.mark.usefixtures("worker_executor_async") async def test_cancelled_legacy_call_finishes_before_returning(cancellations: int) -> None: - started = threading.Event() + loop = asyncio.get_running_loop() + started = asyncio.Event() release = threading.Event() finished = threading.Event() def legacy_operation() -> None: - started.set() - if not release.wait(timeout=10): + loop.call_soon_threadsafe(started.set) + if not release.wait(timeout=TEST_WATCHDOG_SECONDS): raise RuntimeError("Legacy operation was not released") finished.set() task = asyncio.create_task(run_legacy_sync_async(legacy_operation)) try: - assert await asyncio.to_thread(started.wait, 5) + await wait_for_worker_ready_async(started=started, operation=task) for _ in range(cancellations): task.cancel() await asyncio.sleep(0) assert not task.done() finally: release.set() - with pytest.raises(asyncio.CancelledError): - await task + await drain_tasks_async(task) + with pytest.raises(asyncio.CancelledError): + task.result() assert finished.is_set() +@pytest.mark.usefixtures("worker_executor_async") async def test_cancelled_legacy_failure_preserves_cancellation() -> None: - started = threading.Event() + loop = asyncio.get_running_loop() + started = asyncio.Event() release = threading.Event() def fail() -> None: - started.set() - assert release.wait(timeout=10) + loop.call_soon_threadsafe(started.set) + assert release.wait(timeout=TEST_WATCHDOG_SECONDS), "Legacy operation was not released" raise ValueError("worker failed") task = asyncio.create_task(run_legacy_sync_async(fail)) try: - assert await asyncio.to_thread(started.wait, 5) + await wait_for_worker_ready_async(started=started, operation=task) task.cancel() await asyncio.sleep(0) finally: release.set() + await drain_tasks_async(task) with pytest.raises(asyncio.CancelledError) as error: - await task + task.result() assert isinstance(error.value.__cause__, ValueError) @@ -443,12 +471,14 @@ async def test_sync_disposal_rejects_active_async_resources(sqlite_instance: SQL dispose.assert_not_called() +@pytest.mark.usefixtures("worker_executor_async") async def test_sync_read_rejects_same_loop_async_overlap(sqlite_instance: SQLiteMemory) -> None: - started, release = threading.Event(), threading.Event() + loop = asyncio.get_running_loop() + started, release = asyncio.Event(), threading.Event() def delayed_read() -> int: - started.set() - if not release.wait(timeout=5): + loop.call_soon_threadsafe(started.set) + if not release.wait(timeout=TEST_WATCHDOG_SECONDS): raise RuntimeError("Database read was not released") return 1 @@ -461,14 +491,15 @@ def delayed_read() -> int: ) read = asyncio.create_task(session.execute(text("SELECT delayed_read()"))) try: - assert await asyncio.to_thread(started.wait, 5) + await wait_for_worker_ready_async(started=started, operation=read) with pytest.warns(DeprecationWarning), pytest.raises(RuntimeError, match="cannot overlap"): sqlite_instance.get_message_pieces() await asyncio.sleep(0) assert not read.done() finally: release.set() - await read + await drain_tasks_async(read) + read.result() assert await sqlite_instance.get_message_pieces_async() == [] with pytest.warns(DeprecationWarning): assert sqlite_instance.get_message_pieces() == [] @@ -482,11 +513,22 @@ async def test_async_session_rejects_same_thread_sync_overlap(sqlite_instance: S @pytest.mark.parametrize("reader", ["async", "sync"]) +@pytest.mark.usefixtures("worker_executor_async") async def test_in_memory_transactions_are_serialized_across_loops(sqlite_instance: SQLiteMemory, reader: str) -> None: - started = threading.Event() + loop = asyncio.get_running_loop() + blocked = asyncio.Event() + transaction_lock = sqlite_instance._transaction_lock + assert transaction_lock is not None + + def acquire(*, blocking: bool = True) -> bool: + acquired = transaction_lock.acquire(blocking=False) + if not acquired: + loop.call_soon_threadsafe(blocked.set) + if blocking: + return transaction_lock.acquire() + return acquired async def read_and_write_async() -> int: - started.set() try: async with await sqlite_instance.get_session_async() as session: count = await session.scalar(text("SELECT COUNT(*) FROM LockProbe")) @@ -500,7 +542,6 @@ async def read_and_write_async() -> int: def read_and_write() -> int: if reader == "async": return asyncio.run(read_and_write_async()) - started.set() with sqlite_instance._get_sync_session() as session: count = session.scalar(text("SELECT COUNT(*) FROM LockProbe")) assert isinstance(count, int) @@ -508,18 +549,25 @@ def read_and_write() -> int: session.commit() return count - async with await sqlite_instance.get_session_async() as session: - await session.execute(text("CREATE TABLE LockProbe (value INTEGER)")) - await session.commit() - await session.execute(text("INSERT INTO LockProbe VALUES (1)")) - task = asyncio.create_task(asyncio.to_thread(read_and_write)) + task = None + with patch.object( + sqlite_instance, "_transaction_lock", spec=transaction_lock, wraps=transaction_lock + ) as observed_lock: + observed_lock.acquire.side_effect = acquire try: - assert await asyncio.to_thread(started.wait, 5) - await asyncio.sleep(0.05) - assert not task.done() + async with await sqlite_instance.get_session_async() as session: + await session.execute(text("CREATE TABLE LockProbe (value INTEGER)")) + await session.commit() + await session.execute(text("INSERT INTO LockProbe VALUES (1)")) + task = asyncio.create_task(asyncio.to_thread(read_and_write)) + try: + await wait_for_worker_ready_async(started=blocked, operation=task) + assert not task.done() + finally: + await session.commit() finally: - await session.commit() - assert await asyncio.wait_for(task, timeout=5) == 1 + await drain_tasks_async(*get_defined_tasks(task)) + assert task.result() == 1 async with await sqlite_instance.get_session_async() as session: assert await session.scalar(text("SELECT COUNT(*) FROM LockProbe")) == 2 @@ -554,10 +602,13 @@ async def write_async() -> None: await session.commit() task = asyncio.create_task(write_async()) - await asyncio.wait_for(inserted.wait(), timeout=5) - task.cancel() + try: + await wait_for_worker_ready_async(started=inserted, operation=task) + finally: + task.cancel() + await drain_tasks_async(task) with pytest.raises(asyncio.CancelledError): - await task + task.result() async with await sqlite_instance.get_session_async() as session: assert await session.scalar(text("SELECT COUNT(*) FROM CancellationProbe")) == 0 await session.execute(text("INSERT INTO CancellationProbe VALUES (2)")) diff --git a/tests/unit/test_async_utils.py b/tests/unit/test_async_utils.py index e8ea007aaf..5694f577c7 100644 --- a/tests/unit/test_async_utils.py +++ b/tests/unit/test_async_utils.py @@ -3,6 +3,9 @@ import asyncio import inspect +import threading +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext from typing import assert_type from unittest.mock import MagicMock, patch @@ -10,10 +13,12 @@ from unit.async_utils import ( TEST_WATCHDOG_SECONDS, + drain_tasks_async, get_defined_tasks, settle_tasks_async, wait_for_completion_async, wait_for_event_async, + wait_for_worker_ready_async, ) @@ -81,6 +86,153 @@ async def test_wait_for_completion_timeout_does_not_cancel_task_async() -> None: await task +async def test_drain_tasks_preserves_results_and_exceptions_async() -> None: + async def result_async() -> int: + return 42 + + async def fail_async() -> int: + raise ValueError("worker failed") + + result, failure = await drain_tasks_async(asyncio.create_task(result_async()), asyncio.create_task(fail_async())) + assert result == 42 + assert isinstance(failure, ValueError) + assert str(failure) == "worker failed" + + +@pytest.mark.parametrize("cancel_before_cleanup", [False, True]) +async def test_drain_tasks_retains_cancellation_until_worker_finishes_async(cancel_before_cleanup: bool) -> None: + release = asyncio.Event() + draining = asyncio.Event() + worker = asyncio.create_task(release.wait()) + + async def owner_async() -> None: + try: + if cancel_before_cleanup: + raise asyncio.CancelledError("original cancellation") + finally: + draining.set() + await drain_tasks_async(worker) + + owner = asyncio.create_task(owner_async()) + try: + await wait_for_event_async(event=draining, tasks=[owner]) + owner.cancel("cleanup cancellation") + await asyncio.sleep(0) + owner.cancel("repeated cleanup cancellation") + await asyncio.sleep(0) + assert not owner.done() + assert not worker.done() + assert worker.cancelling() == 0 + release.set() + message = "original cancellation" if cancel_before_cleanup else "cleanup cancellation" + with pytest.raises(asyncio.CancelledError, match=f"^{message}$"): + await wait_for_completion_async(future=owner) + assert worker.result() is True + finally: + release.set() + await drain_tasks_async(owner, worker) + + +async def test_drain_tasks_timeout_reports_unfinished_worker_without_cancelling_async() -> None: + release = asyncio.Event() + worker = asyncio.create_task(release.wait()) + try: + with pytest.raises(TimeoutError, match="cleanup watchdog expired"): + await drain_tasks_async(worker, timeout=0) + assert not worker.done() + assert worker.cancelling() == 0 + finally: + release.set() + await drain_tasks_async(worker) + + +@pytest.mark.parametrize("timeout", [-1, float("inf"), float("nan")]) +async def test_drain_tasks_rejects_invalid_watchdog_async(timeout: float) -> None: + with pytest.raises(ValueError, match="finite and non-negative"): + await drain_tasks_async(timeout=timeout) + + +@pytest.mark.parametrize("max_workers", [1, 2], ids=["one-worker", "two-workers"]) +@pytest.mark.parametrize( + "interruption", [None, AssertionError, asyncio.CancelledError], ids=["completion", "assertion", "cancellation"] +) +async def test_worker_readiness_does_not_need_an_executor_worker_async( + *, max_workers: int, interruption: type[BaseException] | None +) -> None: + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = threading.Event() + finished = threading.Event() + + def work() -> int: + loop.call_soon_threadsafe(started.set) + try: + assert release.wait(timeout=TEST_WATCHDOG_SECONDS), "Worker was not released" + return 42 + finally: + finished.set() + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + loop.set_default_executor(executor) + operation = asyncio.create_task(asyncio.to_thread(work)) + outcome = pytest.raises(interruption, match="Interrupted while held") if interruption else nullcontext() + with outcome: + try: + await wait_for_worker_ready_async(started=started, operation=operation) + assert not operation.done() + assert not finished.is_set() + if interruption: + raise interruption("Interrupted while held") + finally: + release.set() + assert await drain_tasks_async(operation) == [42] + assert finished.is_set() + assert operation.done() + + +@pytest.mark.parametrize("outcome", ["success", "failure", "cancellation"]) +async def test_worker_readiness_rejects_completion_after_signal_async(outcome: str) -> None: + started = asyncio.Event() + started.set() + operation: asyncio.Future[int] = asyncio.get_running_loop().create_future() + failure = ValueError("worker failed before readiness") + if outcome == "failure": + operation.set_exception(failure) + expected = pytest.raises(ValueError, match="worker failed before readiness") + elif outcome == "cancellation": + operation.cancel("worker cancelled before readiness") + expected = pytest.raises(asyncio.CancelledError, match="worker cancelled before readiness") + else: + operation.set_result(42) + expected = pytest.raises(AssertionError, match="completed before.*readiness") + + before = asyncio.all_tasks() + with expected as error: + await wait_for_worker_ready_async(started=started, operation=operation, timeout=0) + if outcome == "failure": + assert error.value is failure + assert asyncio.all_tasks() == before + + +@pytest.mark.parametrize("max_workers", [1, 2], ids=["one-worker", "two-workers"]) +async def test_worker_readiness_propagates_thread_failure_before_signal_async(max_workers: int) -> None: + failure = ValueError("worker failed before signaling") + + def fail() -> None: + raise failure + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + asyncio.get_running_loop().set_default_executor(executor) + operation = asyncio.create_task(asyncio.to_thread(fail)) + try: + with pytest.raises(ValueError, match="worker failed before signaling") as error: + await wait_for_worker_ready_async(started=asyncio.Event(), operation=operation) + assert error.value is failure + assert operation.done() + finally: + await drain_tasks_async(operation) + + async def test_wait_for_event_observes_checkpoint_before_completion_async() -> None: reached = asyncio.Event() release = asyncio.Event()