From da37163debdb589edd7f7ecef11d229cab01250f Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Fri, 16 Oct 2026 20:10:16 -0700 Subject: [PATCH 1/2] 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()) == [] From 96f6bfe969a2a83e2ebbe21ba1d20cfb3f1edfb1 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Wed, 21 Oct 2026 09:38:51 -0700 Subject: [PATCH 2/2] TEST Drain owned tasks and verify cancellation ordering Handle teardown watchdog expiry before releasing mocks and fixtures, cover shielded media write ordering, and replace incidental timing checks with explicit readiness synchronization. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: abe68236-1370-4206-b69d-2a72e999262c --- doc/contributing/5_unit_tests.md | 9 ++ tests/unit/async_utils.py | 48 ++++++- tests/unit/backend/test_attack_service.py | 17 +-- .../unit/backend/test_conversation_editor.py | 75 +++++++---- tests/unit/backend/test_media_persistence.py | 85 ++++++++++++- .../unit/backend/test_message_send_service.py | 37 ++++-- tests/unit/test_async_utils.py | 119 +++++++++++++++++- 7 files changed, 348 insertions(+), 42 deletions(-) diff --git a/doc/contributing/5_unit_tests.md b/doc/contributing/5_unit_tests.md index 72f369e473..d1ae269791 100644 --- a/doc/contributing/5_unit_tests.md +++ b/doc/contributing/5_unit_tests.md @@ -29,6 +29,15 @@ When observing an operation's cancellation or cleanup, use `wait_for_completion_ `unit.async_utils`. Unlike `asyncio.wait_for`, its watchdog does not send another cancellation request to the operation when the wait expires. Release blocked workers and drain owned tasks in `finally` so a failed assertion does not leave background work behind. +Use `settle_tasks_async` for the final join: it retains the gather, cancels pending work +if its watchdog expires, and drains cancellation cleanup before raising `TimeoutError`. +That final drain intentionally waits for cooperative cleanup rather than abandoning a task +at a second deadline. An outer test timeout can interrupt this drain; it is a deadlock +backstop, not a guarantee that an uncooperative task settles. A whole-test timeout equal +to an inner watchdog can fire before that watchdog's expiry handling runs. +Use `wait_for_event_async` to watch readiness alongside owned tasks, so a task failure +before readiness is reported immediately instead of waiting for an event. Once readiness +is set, the caller observes the task's outcome separately, even if it has already failed. For deadline tests, expire a real `asyncio.Timeout` with `reschedule` once the operation reaches the intended pending await. Check the configured timeout arguments, cancellation, cleanup, and diff --git a/tests/unit/async_utils.py b/tests/unit/async_utils.py index c4841aa1c8..f7135ecef4 100644 --- a/tests/unit/async_utils.py +++ b/tests/unit/async_utils.py @@ -2,7 +2,9 @@ # Licensed under the MIT license. import asyncio -from typing import TypeVar +from collections.abc import Sequence +from contextlib import suppress +from typing import Any, TypeVar T = TypeVar("T") @@ -26,3 +28,47 @@ async def wait_for_completion_async(*, future: asyncio.Future[T], timeout: float if not done: raise TimeoutError("The operation under test did not complete before the test watchdog expired.") return future.result() + + +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. + + Release test-controlled gates first. Unlike an observation wait, teardown must retain + ownership until cancellation cleanup finishes, even after the watchdog has expired. + """ + joined = asyncio.gather(*tasks, return_exceptions=True) + try: + await wait_for_completion_async(future=joined, timeout=timeout) + except TimeoutError: + joined.cancel() + with suppress(asyncio.CancelledError): + await joined + raise + + +async def wait_for_event_async( + *, + event: asyncio.Event, + tasks: Sequence[asyncio.Future[Any]], + timeout: float = TEST_WATCHDOG_SECONDS, +) -> None: + """Report task completion or failure only when readiness has not been signalled. + + Once readiness is set, the caller owns observing the operation's outcome, including + failures that complete before this readiness observer resumes. + """ + ready = asyncio.create_task(event.wait()) + try: + done, _ = await asyncio.wait({ready, *tasks}, timeout=timeout, return_when=asyncio.FIRST_COMPLETED) + if event.is_set(): + return + if not done: + raise TimeoutError("The operation did not signal readiness before the test watchdog expired.") + for task in tasks: + if task.done(): + task.result() + raise AssertionError("The operation completed without signalling readiness.") + finally: + ready.cancel() + with suppress(asyncio.CancelledError): + await ready diff --git a/tests/unit/backend/test_attack_service.py b/tests/unit/backend/test_attack_service.py index f54b6e2908..01abb574db 100644 --- a/tests/unit/backend/test_attack_service.py +++ b/tests/unit/backend/test_attack_service.py @@ -77,6 +77,7 @@ ) from pyrit.models.conversation_stats import ConversationStats from pyrit.prompt_normalizer import PromptNormalizer +from unit.async_utils import TEST_WATCHDOG_SECONDS, settle_tasks_async, wait_for_event_async from unit.backend.mocks import _make_matching_target_mock, make_attack_result, make_mock_memory from unit.mocks import MockPromptTarget @@ -1569,7 +1570,7 @@ async def test_add_message_propagates_service_errors( await attack_service.add_message_async(attack_result_id="attack", request=request) reader.assert_not_called() - @pytest.mark.timeout(10) + @pytest.mark.timeout(TEST_WATCHDOG_SECONDS) @pytest.mark.parametrize("read_method", ["get_attack_async", "get_conversation_messages_async"]) @pytest.mark.parametrize("incoming_send", [False, True]) async def test_response_assembly_keeps_conversation_reserved_async( @@ -1605,7 +1606,7 @@ async def pause_after_read_async(**kwargs: Any) -> AttackSummary | ConversationM service.add_message_async(attack_result_id=ar.attack_result_id, request=request) ) try: - await started.wait() + await wait_for_event_async(event=started, tasks=[active]) with pytest.raises(ManualSendConflictError): await peer.add_message_async( attack_result_id=ar.attack_result_id, @@ -1618,7 +1619,8 @@ async def pause_after_read_async(**kwargs: Any) -> AttackSummary | ConversationM ) finally: release.set() - result = await active + await settle_tasks_async(active) + result = active.result() assert result.attack.message_count == 2 assert [message.message_pieces[0].original_value for message in result.messages.messages] == [ @@ -1637,7 +1639,7 @@ async def pause_after_read_async(**kwargs: Any) -> AttackSummary | ConversationM ), ) - @pytest.mark.timeout(10) + @pytest.mark.timeout(TEST_WATCHDOG_SECONDS) @pytest.mark.parametrize("read_method", ["get_attack_async", "get_conversation_messages_async"]) @pytest.mark.parametrize("cancel", [False, True]) async def test_response_failure_releases_conversation_async( @@ -1662,14 +1664,15 @@ async def fail_read_async(**_: Any) -> None: ): active = asyncio.create_task(attack_service.add_message_async(attack_result_id="attack", request=request)) try: - await started.wait() + await wait_for_event_async(event=started, tasks=[active]) assert sender._scheduler._conversations == {"conversation"} if cancel: active.cancel() finally: release.set() - with pytest.raises(asyncio.CancelledError if cancel else RuntimeError): - await active + await settle_tasks_async(active) + with pytest.raises(asyncio.CancelledError if cancel else RuntimeError): + active.result() assert not sender._scheduler._conversations await sender.add_message_async(attack_result_id="attack", request=request) diff --git a/tests/unit/backend/test_conversation_editor.py b/tests/unit/backend/test_conversation_editor.py index be702d3780..d0750910b3 100644 --- a/tests/unit/backend/test_conversation_editor.py +++ b/tests/unit/backend/test_conversation_editor.py @@ -46,7 +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.async_utils import TEST_WATCHDOG_SECONDS, settle_tasks_async, wait_for_completion_async, wait_for_event_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 @@ -985,6 +985,49 @@ async def test_media_rollback_deletes_only_staged_files_async(self, sqlite_insta assert not list(Path(sqlite_instance.results_path).rglob("*.txt")) assert (await sqlite_instance.get_attack_results_async()) == [] + @pytest.mark.timeout(TEST_WATCHDOG_SECONDS) + async def test_media_cleanup_watchdog_settles_save_before_teardown_async( + self, *, sqlite_instance: SQLiteMemory + ) -> None: + request = draft() + request.messages[0].pieces = [ + ConversationPieceRequest(data_type="binary_path", original_value="data:text/plain;base64,bWVkaWE=") + ] + storage = sqlite_instance.results_storage_io + write, delete = storage.write_file_async, storage.delete_file_async + cleanup_started, cleanup_finished = asyncio.Event(), asyncio.Event() + + async def failed_write_async(*args: Any, **kwargs: Any) -> None: + await write(*args, **kwargs) + raise OSError("write interrupted") + + async def blocked_delete_async(path: str) -> None: + try: + cleanup_started.set() + await asyncio.Event().wait() + finally: + await delete(path) + cleanup_finished.set() + + with ( + patch.object(storage, "write_file_async", side_effect=failed_write_async), + patch.object(storage, "delete_file_async", side_effect=blocked_delete_async) as delete_file, + ): + task = asyncio.create_task(AttackService().save_conversation_async(request=request)) + try: + await wait_for_event_async(event=cleanup_started, tasks=[task]) + with pytest.raises(TimeoutError, match="test watchdog expired"): + await settle_tasks_async(task, timeout=0) + assert task.cancelled() + assert cleanup_finished.is_set() + delete_file.assert_awaited_once() + finally: + if not task.done(): + task.cancel() + await settle_tasks_async(task) + assert not list(Path(sqlite_instance.results_path).rglob("*.txt")) + assert (await sqlite_instance.get_attack_results_async()) == [] + async def test_message_only_save_keeps_current_objective_async(self, *, sqlite_instance: SQLiteMemory) -> None: service = AttackService() first = await service.save_conversation_async(request=draft()) @@ -1269,18 +1312,21 @@ async def test_transaction_failure_rolls_back_attack_and_conversation_async( assert (await sqlite_instance.get_conversation_metadata_async(conversation_id=str(request.save_id))) is None assert (await sqlite_instance.get_attack_results_async()) == [] + @pytest.mark.timeout(TEST_WATCHDOG_SECONDS) async def test_concurrent_retries_leave_one_conversation_async(self, sqlite_instance: SQLiteMemory) -> None: service = AttackService() request = draft() - first, second = await asyncio.gather( - service.save_conversation_async(request=request), - service.save_conversation_async(request=request), - ) + tasks = [asyncio.create_task(service.save_conversation_async(request=request)) for _ in range(2)] + try: + first, second = await wait_for_completion_async(future=asyncio.gather(*tasks)) + finally: + await settle_tasks_async(*tasks) assert first.messages == second.messages assert len(await sqlite_instance.get_attack_results_async()) == 1 assert len(await sqlite_instance.get_message_pieces_async(conversation_id=str(request.save_id))) == 1 @pytest.mark.parametrize("cancel", [False, True]) + @pytest.mark.timeout(TEST_WATCHDOG_SECONDS) async def test_failed_media_write_cleans_partial_file_async( self, *, sqlite_instance: SQLiteMemory, cancel: bool ) -> None: @@ -1306,31 +1352,18 @@ 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)) - 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.") + await wait_for_event_async(event=started, tasks=[task]) 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(): + if not task.done() and not task.cancelling(): task.cancel() # The shielded writer must finish before the save task can clean up. finish.set() - started_task.cancel() - await wait_for_completion_async( - future=asyncio.gather(task, started_task, return_exceptions=True), - ) + await settle_tasks_async(task) assert not list(Path(sqlite_instance.results_path).rglob("*.txt")) assert (await sqlite_instance.get_attack_results_async()) == [] diff --git a/tests/unit/backend/test_media_persistence.py b/tests/unit/backend/test_media_persistence.py index d44da4ccb0..713f80f019 100644 --- a/tests/unit/backend/test_media_persistence.py +++ b/tests/unit/backend/test_media_persistence.py @@ -3,13 +3,22 @@ """Tests for shared backend media persistence.""" +import asyncio +import uuid from pathlib import Path +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch from urllib.parse import urlencode import pytest -from pyrit.backend.models.attacks import MessagePieceRequest +from pyrit.backend.models.attacks import ( + ConversationMessageRequest, + ConversationPieceRequest, + MessagePieceRequest, + SaveConversationRequest, +) +from pyrit.backend.services.attack_service import AttackService from pyrit.backend.services.media_persistence import ( MediaAccessDeniedError, MediaOrigin, @@ -20,6 +29,7 @@ validate_media_path, ) from pyrit.memory import AzureBlobStorageIO, CentralMemory, DiskStorageIO, SQLiteMemory, data_serializer_factory +from unit.async_utils import TEST_WATCHDOG_SECONDS, settle_tasks_async, wait_for_completion_async, wait_for_event_async def _serializer(*, value: str = "/saved/media.bin") -> MagicMock: @@ -480,3 +490,76 @@ async def test_persistence_failure_returns_no_partial_result() -> None: data_type="binary_path", serializer_factory=MagicMock(return_value=serializer), ) + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.timeout(TEST_WATCHDOG_SECONDS) +async def test_cancelled_save_finishes_writer_before_deleting_media_async(sqlite_instance: SQLiteMemory) -> None: + request = SaveConversationRequest( + save_id=uuid.uuid4(), + destination="new_attack", + objective="Save media", + operator="owner", + messages=[ + ConversationMessageRequest( + role="user", + pieces=[ + ConversationPieceRequest(data_type="binary_path", original_value="data:text/plain;base64,bWVkaWE=") + ], + ) + ], + ) + started, release_writer, writer_finished, cancellation_observed = ( + asyncio.Event(), + asyncio.Event(), + asyncio.Event(), + asyncio.Event(), + ) + storage = sqlite_instance.results_storage_io + write, delete, shield = storage.write_file_async, storage.delete_file_async, asyncio.shield + writers: list[asyncio.Task[Any]] = [] + + async def blocked_write_async(*args: Any, **kwargs: Any) -> None: + writer = asyncio.current_task() + assert writer is not None + writers.append(writer) + started.set() + await release_writer.wait() + await write(*args, **kwargs) + writer_finished.set() + + async def checked_delete_async(path: str) -> None: + assert writer_finished.is_set() + await delete(path) + + def observe_shield(future: asyncio.Future[Any]) -> asyncio.Future[Any]: + protected = shield(future) + protected.add_done_callback(lambda done: cancellation_observed.set() if done.cancelled() else None) + return protected + + with ( + patch.object(storage, "write_file_async", side_effect=blocked_write_async), + patch.object(storage, "delete_file_async", side_effect=checked_delete_async) as delete_file, + patch.object(asyncio, "shield", side_effect=observe_shield), + ): + task = asyncio.create_task(AttackService().save_conversation_async(request=request)) + try: + await wait_for_event_async(event=started, tasks=[task]) + task.cancel() + # Observe delivery to the shield, not merely issuance of a cancellation request. + await wait_for_event_async(event=cancellation_observed, tasks=[task]) + assert not task.done() + assert not writers[0].done() + assert not writer_finished.is_set() + delete_file.assert_not_awaited() + release_writer.set() + with pytest.raises(asyncio.CancelledError): + await wait_for_completion_async(future=task) + assert writer_finished.is_set() + delete_file.assert_awaited_once() + finally: + release_writer.set() + if not task.done() and not task.cancelling(): + task.cancel() + await settle_tasks_async(task, *writers) + assert not list(Path(sqlite_instance.results_path).rglob("*.txt")) diff --git a/tests/unit/backend/test_message_send_service.py b/tests/unit/backend/test_message_send_service.py index d470c2afb7..3ef111bfcc 100644 --- a/tests/unit/backend/test_message_send_service.py +++ b/tests/unit/backend/test_message_send_service.py @@ -63,7 +63,12 @@ from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target.common.target_capabilities import CapabilityName, TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from unit.async_utils import wait_for_completion_async +from unit.async_utils import ( + TEST_WATCHDOG_SECONDS, + settle_tasks_async, + wait_for_completion_async, + wait_for_event_async, +) from unit.backend.mocks import ( _make_matching_target_mock, _settle_send_async, @@ -1959,6 +1964,7 @@ async def test_dispatch_failure_releases_ownership_without_reusing_old_errors_as send.side_effect = None await message_send_service.add_message_async(attack_result_id="attack", request=_request()) + @pytest.mark.timeout(TEST_WATCHDOG_SECONDS) @pytest.mark.parametrize("write_method", ["add_message_pieces_to_memory_async", "update_attack_result_by_id_async"]) async def test_cancellation_joins_memory_write_before_releasing_ownership_async( self, @@ -1972,7 +1978,7 @@ async def test_cancellation_joins_memory_write_before_releasing_ownership_async( async def hold_write_async(**_: Any) -> None: started.set() - await asyncio.wait_for(release.wait(), timeout=5) + await release.wait() finished.set() with patch.object(mock_memory, write_method, side_effect=hold_write_async): @@ -1980,7 +1986,7 @@ async def hold_write_async(**_: Any) -> None: message_send_service.add_message_async(attack_result_id="attack", request=_request(send=False)) ) try: - await started.wait() + await wait_for_event_async(event=started, tasks=[active]) for _ in range(2): active.cancel() await asyncio.sleep(0) @@ -1993,8 +1999,9 @@ async def hold_write_async(**_: Any) -> None: ) finally: release.set() - with pytest.raises(asyncio.CancelledError): - await active + await settle_tasks_async(active) + with pytest.raises(asyncio.CancelledError): + active.result() assert finished.is_set() assert not message_send_service._scheduler._conversations assert not message_send_service._scheduler._metadata_updates @@ -2097,6 +2104,7 @@ async def hold_first_async(*, conversation_id: str, **_: Any) -> None: assert stored_target is not None assert stored_target.hash == target.get_identifier().hash + @pytest.mark.timeout(TEST_WATCHDOG_SECONDS) @pytest.mark.parametrize("preconverted", [False, True]) async def test_parallel_sends_serialize_complete_metadata_updates_async( self, @@ -2123,6 +2131,7 @@ async def test_parallel_sends_serialize_complete_metadata_updates_async( second_target_started, first_write_started = asyncio.Event(), asyncio.Event() second_metadata_attempted, release_first_write = asyncio.Event(), asyncio.Event() + second_metadata_acquired = asyncio.Event() writes: list[dict[str, Any]] = [] attempts = 0 send = target._send_prompt_to_target_async @@ -2142,16 +2151,19 @@ async def overlap_targets_async(*, normalized_conversation: list[Message]) -> li async def observe_metadata_async(*, attack_result_id: str) -> AsyncGenerator[None, None]: nonlocal attempts attempts += 1 - if attempts == 2: + second_attempt = attempts == 2 + if second_attempt: second_metadata_attempted.set() async with guard(attack_result_id=attack_result_id): + if second_attempt: + second_metadata_acquired.set() yield async def delay_first_write_async(*, attack_result_id: str, update_fields: dict[str, Any]) -> None: writes.append(update_fields) if len(writes) == 1: first_write_started.set() - await asyncio.wait_for(release_first_write.wait(), timeout=5) + await release_first_write.wait() await update(attack_result_id=attack_result_id, update_fields=update_fields) with ( @@ -2168,15 +2180,18 @@ async def delay_first_write_async(*, attack_result_id: str, update_fields: dict[ for service, request in zip([first, second], requests, strict=True) ] try: - await second_metadata_attempted.wait() + await wait_for_event_async(event=second_metadata_attempted, tasks=tasks) assert all(not task.done() for task in tasks) - with pytest.raises(TimeoutError): - await asyncio.wait_for(asyncio.shield(tasks[1]), timeout=0.1) + assert not second_metadata_acquired.is_set() assert len(writes) == 1 assert not tasks[1].done() finally: + second_target_started.set() + first_write_started.set() release_first_write.set() - await asyncio.gather(*tasks) + await settle_tasks_async(*tasks) + for task in tasks: + task.result() assert len(writes) == 2 assert writes[0]["timestamp"] <= writes[1]["timestamp"] diff --git a/tests/unit/test_async_utils.py b/tests/unit/test_async_utils.py index 474631c369..0fcbce2718 100644 --- a/tests/unit/test_async_utils.py +++ b/tests/unit/test_async_utils.py @@ -4,10 +4,17 @@ import asyncio import inspect from typing import assert_type +from unittest.mock import MagicMock, patch import pytest -from unit.async_utils import get_defined_tasks, wait_for_completion_async +from unit.async_utils import ( + TEST_WATCHDOG_SECONDS, + get_defined_tasks, + settle_tasks_async, + wait_for_completion_async, + wait_for_event_async, +) def test_get_defined_tasks_before_task_creation() -> None: @@ -72,3 +79,113 @@ async def test_wait_for_completion_timeout_does_not_cancel_task_async() -> None: finally: release.set() await task + + +async def test_settle_tasks_joins_peers_after_failure_async() -> None: + release = asyncio.Event() + finished = asyncio.Event() + + async def fail_async() -> None: + raise RuntimeError("operation failed") + + async def peer_async() -> None: + await release.wait() + finished.set() + + failed, peer = asyncio.create_task(fail_async()), asyncio.create_task(peer_async()) + release.set() + await settle_tasks_async(failed, peer) + assert failed.done() and peer.done() + assert finished.is_set() + + +@pytest.mark.timeout(TEST_WATCHDOG_SECONDS) +async def test_settle_tasks_watchdog_drains_cleanup_before_patch_teardown_async() -> None: + started, cleanup_started, release_cleanup = asyncio.Event(), asyncio.Event(), asyncio.Event() + resource = MagicMock() + + async def blocked_async() -> None: + try: + started.set() + await asyncio.Event().wait() + finally: + cleanup_started.set() + await release_cleanup.wait() + resource.close() + + task = asyncio.create_task(blocked_async()) + settlement: asyncio.Task[None] | None = None + with patch.object(resource, "close") as close: + try: + await wait_for_event_async(event=started, tasks=[task]) + settlement = asyncio.create_task(settle_tasks_async(task, timeout=0)) + await wait_for_event_async(event=cleanup_started, tasks=[settlement]) + assert not task.done() + assert not settlement.done() + close.assert_not_called() + release_cleanup.set() + with pytest.raises(TimeoutError, match="test watchdog expired"): + await wait_for_completion_async(future=settlement) + assert task.cancelled() + close.assert_called_once() + finally: + release_cleanup.set() + task.cancel() + await settle_tasks_async(*get_defined_tasks(task, settlement)) + + +async def test_wait_for_event_reports_readiness_async() -> None: + event = asyncio.Event() + event.set() + await wait_for_event_async(event=event, tasks=[]) + + +async def test_wait_for_event_preserves_failure_before_readiness_async() -> None: + event = asyncio.Event() + future: asyncio.Future[None] = asyncio.get_running_loop().create_future() + future.set_exception(RuntimeError("early failure")) + with pytest.raises(RuntimeError, match="early failure"): + await wait_for_event_async(event=event, tasks=[future]) + + +@pytest.mark.parametrize("complete_before_wait", [False, True]) +async def test_wait_for_event_readiness_leaves_completed_failure_for_caller_async( + *, complete_before_wait: bool +) -> None: + event = asyncio.Event() + failure = OSError("operation failed after readiness") + + async def signal_and_fail_async() -> None: + event.set() + raise failure + + task = asyncio.create_task(signal_and_fail_async()) + try: + if complete_before_wait: + await asyncio.wait({task}) + await wait_for_event_async(event=event, tasks=[task]) + assert task.done() + with pytest.raises(OSError, match="failed after readiness") as error: + await wait_for_completion_async(future=task) + assert error.value is failure + finally: + await settle_tasks_async(task) + + +async def test_wait_for_event_rejects_completion_without_readiness_async() -> None: + future: asyncio.Future[None] = asyncio.get_running_loop().create_future() + future.set_result(None) + with pytest.raises(AssertionError, match="without signalling readiness"): + await wait_for_event_async(event=asyncio.Event(), tasks=[future]) + + +async def test_wait_for_event_timeout_does_not_cancel_operation_async() -> None: + task = asyncio.create_task(asyncio.Event().wait()) + try: + with pytest.raises(TimeoutError, match="test watchdog expired"): + await wait_for_event_async(event=asyncio.Event(), tasks=[task], timeout=0) + assert not task.done() + assert not task.cancelling() + finally: + task.cancel() + await settle_tasks_async(task)