From b9bdcc528e532c0de241adc3e0763ef6eba7fc0a Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Tue, 6 Oct 2026 12:50:13 -0700 Subject: [PATCH 1/3] TEST Stabilize merge-queue async and browser checks Expire generator deadlines at the intended await, observe cancellation without watchdog-induced cancellation, and allow CI scheduling slack for coordination. Remove native spinner press-and-hold from single-step browser checks and document the async test patterns. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- doc/contributing/5_unit_tests.md | 15 +++++ frontend/e2e/numeric-controls.spec.ts | 2 +- tests/unit/async_utils.py | 15 +++++ tests/unit/backend/test_main.py | 6 +- tests/unit/backend/test_scenario_service.py | 38 +++++++---- .../test_target_objective_generator.py | 65 ++++++++++++++----- .../core/test_scenario_partial_results.py | 27 ++++---- tests/unit/test_async_utils.py | 49 ++++++++++++++ 8 files changed, 168 insertions(+), 49 deletions(-) create mode 100644 tests/unit/async_utils.py create mode 100644 tests/unit/test_async_utils.py diff --git a/doc/contributing/5_unit_tests.md b/doc/contributing/5_unit_tests.md index 0132aa12c2..5adaeeaa9e 100644 --- a/doc/contributing/5_unit_tests.md +++ b/doc/contributing/5_unit_tests.md @@ -15,6 +15,21 @@ Testing is an art to get right! But here are some best practices in terms of uni Not all of our current tests follow these practices (we're working on it!) But for some good examples, see [test_tts_send_prompt_file_save_async](../../tests/unit/prompt_target/target/test_tts_target.py), which has many of these best practices incorporated in the test. +## Async timing and cancellation + +Use events to coordinate concurrent operations and assert their ordering or concurrency bounds. +Timeouts that only prevent a test from hanging should allow for busy CI runners, rather than +acting as performance assertions. + +When observing an operation's cancellation or cleanup, use `wait_for_completion_async` from +`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. + +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 +original outcome. This avoids short wall-clock deadlines expiring during unrelated setup. + ## SQLite memory fixtures `sqlite_instance` stays function-scoped. Each test gets a fresh in-memory database and results directory, and its SQLite singleton and CentralMemory registrations are restored afterward. The fixture owns disposal of its memory instance instead of registering process-exit cleanup callbacks for every test. diff --git a/frontend/e2e/numeric-controls.spec.ts b/frontend/e2e/numeric-controls.spec.ts index 58558548ba..319d8cb74f 100644 --- a/frontend/e2e/numeric-controls.spec.ts +++ b/frontend/e2e/numeric-controls.spec.ts @@ -79,8 +79,8 @@ async function clickNativeSpinner(input: Locator, direction: 'up' | 'down'): Pro if (!box) throw new Error('Expected a visible numeric input.') const paddingRight = await input.evaluate((element) => parseFloat(getComputedStyle(element).paddingRight)) // Chromium's native spinner is a UA shadow control, not an accessible button. + // Native spinners auto-repeat when held, so a single-step check must not hold the mouse down. await input.click({ - delay: 200, position: { x: box.width - paddingRight - 8, y: box.height / 2 + (direction === 'up' ? -4 : 4) }, }) } diff --git a/tests/unit/async_utils.py b/tests/unit/async_utils.py new file mode 100644 index 0000000000..4c62c910b8 --- /dev/null +++ b/tests/unit/async_utils.py @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +from typing import TypeVar + +T = TypeVar("T") + + +async def wait_for_completion_async(*, future: asyncio.Future[T], timeout: float = 30) -> T: + """Bound a test wait without injecting cancellation into the operation under test.""" + done, _ = await asyncio.wait({future}, timeout=timeout) + if not done: + raise TimeoutError("The operation under test did not complete before the test watchdog expired.") + return future.result() diff --git a/tests/unit/backend/test_main.py b/tests/unit/backend/test_main.py index fe32bf2a00..0f78a16d0e 100644 --- a/tests/unit/backend/test_main.py +++ b/tests/unit/backend/test_main.py @@ -42,7 +42,7 @@ async def test_health_responds_while_database_operation_is_pending(sqlite_instan def wait_in_database() -> int: started.set() - if not release.wait(timeout=10): + if not release.wait(timeout=60): raise RuntimeError("Database wait was not released") return 1 @@ -55,9 +55,9 @@ def wait_in_database() -> int: ) query = asyncio.create_task(session.execute(text("SELECT wait_in_database()"))) try: - assert await asyncio.to_thread(started.wait, 5) + assert await asyncio.to_thread(started.wait, 30) async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client: - response = await asyncio.wait_for(client.get("/api/health"), timeout=2) + response = await asyncio.wait_for(client.get("/api/health"), timeout=30) assert response.status_code == 200 assert response.json()["status"] == "healthy" assert not query.done() diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py index 4b3551baa3..ee6bb2f21b 100644 --- a/tests/unit/backend/test_scenario_service.py +++ b/tests/unit/backend/test_scenario_service.py @@ -769,16 +769,21 @@ async def estimate_async( service._run_default_estimate_async = AsyncMock(side_effect=estimate_async) catalog_task = asyncio.create_task(service.list_scenarios_async()) - await asyncio.wait_for(two_started.wait(), timeout=2) - await asyncio.sleep(0) + try: + await asyncio.wait_for(two_started.wait(), timeout=30) + await asyncio.sleep(0) - assert service._run_default_estimate_async.await_count == 2 - release.set() - result = await catalog_task + assert service._run_default_estimate_async.await_count == 2 + release.set() + result = await catalog_task - assert maximum_active == 2 - assert service._run_default_estimate_async.await_count == 3 - assert all(item.default_run_size == estimate for item in result.items) + assert maximum_active == 2 + assert service._run_default_estimate_async.await_count == 3 + assert all(item.default_run_size == estimate for item in result.items) + finally: + release.set() + await service.close_async() + await asyncio.gather(catalog_task, return_exceptions=True) async def test_catalog_queue_wait_does_not_start_execution_timeout(self) -> None: """A queued catalog estimate starts its timeout only after acquiring capacity.""" @@ -1259,13 +1264,18 @@ async def estimate_async( ) for index in range(3) ] - await asyncio.wait_for(two_started.wait(), timeout=1) - await asyncio.sleep(0) + try: + await asyncio.wait_for(two_started.wait(), timeout=30) + await asyncio.sleep(0) - assert service._estimate_configured_run_size_async.await_count == 2 - release.set() - assert await asyncio.gather(*tasks) == [estimate, estimate, estimate] - assert maximum_active == 2 + assert service._estimate_configured_run_size_async.await_count == 2 + release.set() + assert await asyncio.gather(*tasks) == [estimate, estimate, estimate] + assert maximum_active == 2 + finally: + release.set() + await service.close_async() + await asyncio.gather(*tasks, return_exceptions=True) async def test_metadata_catalog_remains_responsive_during_estimate(self) -> None: """Metadata-only catalog requests do not wait for running estimates.""" diff --git a/tests/unit/executor/promptgen/test_target_objective_generator.py b/tests/unit/executor/promptgen/test_target_objective_generator.py index c7ef984ff8..46e1cc3dd4 100644 --- a/tests/unit/executor/promptgen/test_target_objective_generator.py +++ b/tests/unit/executor/promptgen/test_target_objective_generator.py @@ -36,6 +36,21 @@ def _isolate_pre_send_io(generator: TargetObjectiveGenerator) -> Iterator[None]: yield +@contextmanager +def _controlled_timeouts() -> Iterator[list[tuple[float | None, asyncio.Timeout]]]: + """Expire real asyncio deadlines at the intended await, not during unrelated setup.""" + timeouts: list[tuple[float | None, asyncio.Timeout]] = [] + original_timeout = asyncio.timeout + + def capture_timeout(delay: float | None) -> asyncio.Timeout: + timeout = original_timeout(None) + timeouts.append((delay, timeout)) + return timeout + + with patch("pyrit.executor.promptgen.target_objective_generator.asyncio.timeout", new=capture_timeout): + yield timeouts + + @pytest.mark.usefixtures("patch_central_database") class TestTargetObjectiveGenerator: async def test_valid_batch_and_evidence_async(self, sqlite_instance: MemoryInterface) -> None: @@ -231,6 +246,7 @@ def test_nontext_system_prompt_rejected(self) -> None: target=MockPromptTarget(), system_prompt=SeedPrompt(value="image.png", data_type="image_path") ) + @pytest.mark.timeout(30) async def test_timeout_bounds_pending_send_and_cleanup_async(self, caplog: pytest.LogCaptureFixture) -> None: generator = TargetObjectiveGenerator(target=MockPromptTarget(), timeout_seconds=1) cancelled = asyncio.Event() @@ -238,33 +254,38 @@ async def test_timeout_bounds_pending_send_and_cleanup_async(self, caplog: pytes async def wait_forever_async(**kwargs: object) -> None: try: + timeouts[0][1].reschedule(asyncio.get_running_loop().time()) await asyncio.Event().wait() finally: cancelled.set() async def reset_async(*, conversation_id: str) -> None: try: + timeouts[-1][1].reschedule(asyncio.get_running_loop().time()) await asyncio.Event().wait() finally: cleanup_cancelled.set() with ( _isolate_pre_send_io(generator), + _controlled_timeouts() as timeouts, patch.object(generator, "_CLEANUP_TIMEOUT_SECONDS", 0.01), patch.object( generator._normalizer, "send_prompt_async", new_callable=AsyncMock, side_effect=wait_forever_async ) as send, patch.object(generator._target, "reset_conversation_async", side_effect=reset_async) as reset, ): - async with asyncio.timeout(3): - with pytest.raises(TimeoutError): - await generator.execute_async(instructions="Test", count=2) + with pytest.raises(TimeoutError): + await generator.execute_async(instructions="Test", count=2) + assert [delay for delay, _ in timeouts] == [1, 0.01] + assert all(timeout.expired() for _, timeout in timeouts) send.assert_awaited_once() assert cancelled.is_set() assert cleanup_cancelled.is_set() reset.assert_awaited_once() assert "Timed out resetting generation conversation" in caplog.text + @pytest.mark.timeout(30) @pytest.mark.parametrize( ("pending_method", "conversation_initialized"), [ @@ -283,21 +304,24 @@ async def test_timeout_before_send_preserves_cleanup_contract_async( async def wait_forever_async(**kwargs: object) -> None: try: + timeouts[0][1].reschedule(asyncio.get_running_loop().time()) await asyncio.Event().wait() finally: cancelled.set() with ( _isolate_pre_send_io(generator), + _controlled_timeouts() as timeouts, patch.object( pending_owner, pending_method, new_callable=AsyncMock, side_effect=wait_forever_async ) as pending, patch.object(generator._normalizer, "send_prompt_async", new_callable=AsyncMock) as send, patch.object(generator._target, "reset_conversation_async", new_callable=AsyncMock) as reset, ): - async with asyncio.timeout(3): - with pytest.raises(TimeoutError): - await generator.execute_with_context_async(context=context) + with pytest.raises(TimeoutError): + await generator.execute_with_context_async(context=context) + assert timeouts[0][0] == 1 + assert timeouts[0][1].expired() pending.assert_awaited_once() assert cancelled.is_set() assert context._used @@ -308,6 +332,7 @@ async def wait_forever_async(**kwargs: object) -> None: else: reset.assert_not_awaited() + @pytest.mark.timeout(30) @pytest.mark.parametrize("failure", [None, ConnectionError("Generation failed"), asyncio.CancelledError()]) async def test_cleanup_timeout_preserves_outcome_async( self, *, failure: BaseException | None, caplog: pytest.LogCaptureFixture @@ -318,12 +343,14 @@ async def test_cleanup_timeout_preserves_outcome_async( async def reset_async(*, conversation_id: str) -> None: try: + timeouts[-1][1].reschedule(asyncio.get_running_loop().time()) await asyncio.Event().wait() finally: cleanup_cancelled.set() with ( _isolate_pre_send_io(generator), + _controlled_timeouts() as timeouts, patch.object(generator, "_CLEANUP_TIMEOUT_SECONDS", 0.01), patch.object( generator._normalizer, @@ -334,18 +361,20 @@ async def reset_async(*, conversation_id: str) -> None: ), patch.object(generator._target, "reset_conversation_async", side_effect=reset_async) as reset, ): - async with asyncio.timeout(1): - if failure is None: - result = await generator.execute_async(instructions="Test", count=1) - assert result.objectives == ["A goal"] - elif isinstance(failure, asyncio.CancelledError): - with pytest.raises(asyncio.CancelledError) as error: - await generator.execute_async(instructions="Test", count=1) - assert error.value is failure - else: - with pytest.raises(RuntimeError) as generation_error: - await generator.execute_async(instructions="Test", count=1) - assert generation_error.value.__cause__ is failure + if failure is None: + result = await generator.execute_async(instructions="Test", count=1) + assert result.objectives == ["A goal"] + elif isinstance(failure, asyncio.CancelledError): + with pytest.raises(asyncio.CancelledError) as error: + await generator.execute_async(instructions="Test", count=1) + assert error.value is failure + else: + with pytest.raises(RuntimeError) as generation_error: + await generator.execute_async(instructions="Test", count=1) + assert generation_error.value.__cause__ is failure + assert [delay for delay, _ in timeouts] == [generator._timeout_seconds, 0.01] + assert not timeouts[0][1].expired() + assert timeouts[1][1].expired() reset.assert_awaited_once() assert cleanup_cancelled.is_set() assert "Timed out resetting generation conversation" in caplog.text diff --git a/tests/unit/scenario/core/test_scenario_partial_results.py b/tests/unit/scenario/core/test_scenario_partial_results.py index 24b0c36deb..79eb0eedb4 100644 --- a/tests/unit/scenario/core/test_scenario_partial_results.py +++ b/tests/unit/scenario/core/test_scenario_partial_results.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest +from unit.async_utils import wait_for_completion_async from pyrit.exceptions import ScenarioPartialFailureException from pyrit.executor.attack import PromptSendingAttack @@ -601,11 +602,11 @@ async def run_async(**_kwargs: object) -> AttackExecutorResult[AttackResult]: parent = asyncio.create_task(scenario.run_async()) try: - await asyncio.wait_for(all_started.wait(), timeout=5) + await asyncio.wait_for(all_started.wait(), timeout=30) release_workers.set() parent.cancel("stop scenario") with pytest.raises(asyncio.CancelledError, match="stop scenario"): - await asyncio.wait_for(parent, timeout=5) + await wait_for_completion_async(future=parent) assert all(worker.done() for worker in worker_tasks) assert not scenario._active_atomic_groups @@ -673,10 +674,10 @@ async def sibling_run_async(**_kwargs: object) -> AttackExecutorResult[AttackRes parent = asyncio.create_task(scenario.run_async()) try: - await asyncio.wait_for(all_started.wait(), timeout=5) + await asyncio.wait_for(all_started.wait(), timeout=30) release_sibling.set() parent.cancel("stop scenario") - await asyncio.wait_for(cleanup_started.wait(), timeout=5) + await asyncio.wait_for(cleanup_started.wait(), timeout=30) assert not parent.done() assert worker_tasks[0].cancelling() == 1 [stored] = await scenario._memory.get_scenario_results_async( @@ -692,7 +693,7 @@ async def sibling_run_async(**_kwargs: object) -> AttackExecutorResult[AttackRes release_cleanup.set() with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(parent, timeout=5) + await wait_for_completion_async(future=parent) assert cleanup_finished.is_set() assert all(worker.done() for worker in worker_tasks) assert not scenario._active_atomic_groups @@ -743,9 +744,9 @@ async def cancel_then_resume_async() -> ScenarioResult: parent = asyncio.create_task(cancel_then_resume_async()) try: - await asyncio.wait_for(started.wait(), timeout=5) + await asyncio.wait_for(started.wait(), timeout=30) parent.cancel("stop first run") - result = await asyncio.wait_for(parent, timeout=5) + result = await wait_for_completion_async(future=parent) assert result.scenario_run_state is ScenarioRunState.COMPLETED assert result.number_tries == 2 assert result.get_objectives() == ["objective"] @@ -787,12 +788,12 @@ async def sibling_run_async(**_kwargs): task = asyncio.create_task(scenario.run_async()) try: - await asyncio.wait_for(cleanup_started.wait(), timeout=5) + await asyncio.wait_for(cleanup_started.wait(), timeout=30) task.cancel("caller cancelled during cleanup") await asyncio.sleep(0) allow_cleanup.set() with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(task, timeout=5) + await wait_for_completion_async(future=task) assert cleanup_finished.is_set() assert not scenario._active_atomic_groups finally: @@ -854,10 +855,10 @@ async def observe_fast_worker_async(**kwargs): ): parent = asyncio.create_task(scenario.run_async()) try: - await asyncio.wait_for(all_started.wait(), timeout=5) + await asyncio.wait_for(all_started.wait(), timeout=30) parent.cancel("stop scenario") - await asyncio.wait_for(cleanup_started.wait(), timeout=5) - await asyncio.wait_for(fast_worker_finished.wait(), timeout=5) + await asyncio.wait_for(cleanup_started.wait(), timeout=30) + await asyncio.wait_for(fast_worker_finished.wait(), timeout=30) assert sends["slow"].cancelling() == 1 assert not cleanup_finished.is_set() assert not parent.done() @@ -876,7 +877,7 @@ async def observe_fast_worker_async(**kwargs): release_cleanup.set() with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(parent, timeout=5) + await wait_for_completion_async(future=parent) assert cleanup_finished.is_set() assert set(sends) == {"slow", "fast"} assert all(task.done() for task in sends.values()) diff --git a/tests/unit/test_async_utils.py b/tests/unit/test_async_utils.py new file mode 100644 index 0000000000..813456661f --- /dev/null +++ b/tests/unit/test_async_utils.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio + +import pytest + +from unit.async_utils import wait_for_completion_async + + +async def test_wait_for_completion_returns_result_async() -> None: + future: asyncio.Future[int] = asyncio.get_running_loop().create_future() + future.set_result(42) + + assert await wait_for_completion_async(future=future) == 42 + + +async def test_wait_for_completion_preserves_exception_async() -> None: + future: asyncio.Future[int] = asyncio.get_running_loop().create_future() + failure = RuntimeError("operation failed") + future.set_exception(failure) + + with pytest.raises(RuntimeError, match="operation failed") as error: + await wait_for_completion_async(future=future) + assert error.value is failure + + +async def test_wait_for_completion_preserves_cancellation_async() -> None: + future: asyncio.Future[int] = asyncio.get_running_loop().create_future() + future.cancel("requested cancellation") + + with pytest.raises(asyncio.CancelledError, match="requested cancellation"): + await wait_for_completion_async(future=future) + + +async def test_wait_for_completion_timeout_does_not_cancel_task_async() -> None: + release = asyncio.Event() + task = asyncio.create_task(release.wait()) + try: + with pytest.raises(TimeoutError, match="test watchdog expired"): + await wait_for_completion_async(future=task, timeout=0) + assert not task.done() + assert task.cancelling() == 0 + + release.set() + assert await wait_for_completion_async(future=task) + finally: + release.set() + await task From 8e20720f7c713ebbedee7f26bd4251e7c44a402c Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Tue, 6 Oct 2026 13:31:30 -0700 Subject: [PATCH 2/3] TEST Stabilize SQLite cancellation coordination watchdogs Allow CI scheduling slack before requesting cancellation and use non-cancelling completion waits while preserving rollback, cleanup, and exception assertions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/unit/memory/test_sqlite_cancellation.py | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/tests/unit/memory/test_sqlite_cancellation.py b/tests/unit/memory/test_sqlite_cancellation.py index 0f41acea0a..f97fe6eed2 100644 --- a/tests/unit/memory/test_sqlite_cancellation.py +++ b/tests/unit/memory/test_sqlite_cancellation.py @@ -22,6 +22,7 @@ from pyrit.memory import sqlite_memory as sqlite_memory_module from pyrit.memory.sqlite_memory import _CursorClosingSQLiteConnection from pyrit.models import AttackOutcome, AttackResult, MessagePiece, MessageScorable, Score +from unit.async_utils import wait_for_completion_async from unit.mocks import get_mock_scorer_identifier @@ -124,9 +125,9 @@ async def delayed_connection_close_async(connection: Connection) -> None: ): task = asyncio.create_task(sqlite_memory_async.add_scores_to_memory_async(scores=[persistable_score_async])) try: - await asyncio.wait_for(started.wait(), timeout=5) + await asyncio.wait_for(started.wait(), timeout=30) task.cancel("cancel score transaction") - await asyncio.wait_for(closing_connection.wait(), timeout=5) + await asyncio.wait_for(closing_connection.wait(), timeout=30) for _ in range(additional_cancellations): task.cancel("cancel connection cleanup again") await asyncio.sleep(0) @@ -134,7 +135,7 @@ async def delayed_connection_close_async(connection: Connection) -> None: assert not connection_closed.is_set() release_connection.set() with pytest.raises(asyncio.CancelledError, match="cancel score transaction") as raised: - await task + await wait_for_completion_async(future=task) finally: release_read.set() release_connection.set() @@ -180,9 +181,9 @@ async def write_async() -> None: with patch.object(Connection, "rollback", new=delayed_rollback_async): task = asyncio.create_task(write_async()) try: - await asyncio.wait_for(inserted.wait(), timeout=5) + await asyncio.wait_for(inserted.wait(), timeout=30) task.cancel("cancel transaction body") - await asyncio.wait_for(rolling_back.wait(), timeout=5) + await asyncio.wait_for(rolling_back.wait(), timeout=30) for _ in range(additional_cancellations): task.cancel("cancel rollback again") await asyncio.sleep(0) @@ -190,7 +191,7 @@ async def write_async() -> None: assert not rolled_back.is_set() release_rollback.set() with pytest.raises(asyncio.CancelledError, match="cancel transaction body"): - await task + await wait_for_completion_async(future=task) finally: release_rollback.set() if not task.done(): @@ -257,10 +258,10 @@ async def write_async() -> None: patch.object(Connection, "rollback", new=failed_rollback_async), patch.object(Connection, "close", new=observed_close_async), ): - await asyncio.wait_for(inserted.wait(), timeout=5) + await asyncio.wait_for(inserted.wait(), timeout=30) if cancel_in_body: task.cancel(cancellation_message) - await asyncio.wait_for(rolling_back.wait(), timeout=5) + await asyncio.wait_for(rolling_back.wait(), timeout=30) if cancel_in_body is False: task.cancel(cancellation_message) await asyncio.sleep(0) @@ -272,7 +273,7 @@ async def write_async() -> None: expected = OperationalError if cancel_in_body is None else asyncio.CancelledError message = "rollback failed" if cancel_in_body is None else cancellation_message with pytest.raises(expected, match=message) as raised: - await task + await wait_for_completion_async(future=task) assert connection_closed.is_set() error = raised.value if cancel_in_body is None else raised.value.__cause__ @@ -301,7 +302,7 @@ async def test_cancelled_database_worker_finishes_before_return_async( def wait_in_database() -> int: started.set() - if not release.wait(timeout=10): + if not release.wait(timeout=60): raise RuntimeError("Database worker was not released") finished.set() return 1 @@ -330,9 +331,9 @@ async def read_async() -> None: with patch.object(sqlite_memory_module, "_finish_sqlite_cleanup_async", new=observed_cleanup_async): task = asyncio.create_task(read_async()) try: - assert await asyncio.to_thread(started.wait, 5) + assert await asyncio.to_thread(started.wait, 30) task.cancel("cancel database worker") - await asyncio.wait_for(closing.wait(), timeout=5) + await asyncio.wait_for(closing.wait(), timeout=30) task.cancel("cancel worker cleanup again") await asyncio.sleep(0) assert not task.done() @@ -340,7 +341,7 @@ async def read_async() -> None: assert not closed.is_set() release.set() with pytest.raises(asyncio.CancelledError, match="cancel database worker"): - await task + await wait_for_completion_async(future=task) finally: release.set() if not task.done(): @@ -405,9 +406,9 @@ async def delayed_session_close_async(session: AsyncSession) -> None: ): task = asyncio.create_task(sqlite_memory_async.add_scores_to_memory_async(scores=[persistable_score_async])) try: - await asyncio.wait_for(started.wait(), timeout=5) + await asyncio.wait_for(started.wait(), timeout=30) task.cancel("cancel failed cleanup") - await asyncio.wait_for(closing_cleanup.wait(), timeout=5) + await asyncio.wait_for(closing_cleanup.wait(), timeout=30) for _ in range(additional_cancellations): task.cancel("cancel cleanup again") await asyncio.sleep(0) @@ -417,7 +418,7 @@ async def delayed_session_close_async(session: AsyncSession) -> None: assert not connection_closed.is_set() release_cleanup.set() with pytest.raises(asyncio.CancelledError) as raised: - await task + await wait_for_completion_async(future=task) finally: release.set() release_cleanup.set() From da52483749e0494f4a10fc739c018a8b1bcb4ba3 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Wed, 7 Oct 2026 10:25:27 -0700 Subject: [PATCH 3/3] TEST Separate target pacing isolation from persistence timing Use explicit pacing and metadata gates to prove independent target progress before finalization. Bound observation without injecting cancellation, release and drain owned tasks on failure, and allow the method watchdog to cover its coordination waits. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- doc/contributing/5_unit_tests.md | 4 +++ .../unit/backend/test_message_send_service.py | 33 ++++++++++++++++--- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/doc/contributing/5_unit_tests.md b/doc/contributing/5_unit_tests.md index 5adaeeaa9e..72f369e473 100644 --- a/doc/contributing/5_unit_tests.md +++ b/doc/contributing/5_unit_tests.md @@ -21,6 +21,10 @@ Use events to coordinate concurrent operations and assert their ordering or conc Timeouts that only prevent a test from hanging should allow for busy CI runners, rather than acting as performance assertions. +For isolation tests, hold one operation at an explicit gate and observe the other operation +reaching its intended milestone before releasing that gate. Wait for unrelated persistence or +finalization separately, rather than including it in a short deadline for the behavior under test. + When observing an operation's cancellation or cleanup, use `wait_for_completion_async` from `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 diff --git a/tests/unit/backend/test_message_send_service.py b/tests/unit/backend/test_message_send_service.py index a201a06d6f..56be583eed 100644 --- a/tests/unit/backend/test_message_send_service.py +++ b/tests/unit/backend/test_message_send_service.py @@ -57,6 +57,7 @@ PromptDataType, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer +from unit.async_utils import wait_for_completion_async from unit.backend.mocks import _make_matching_target_mock, make_attack_result, make_mock_memory from unit.mocks import MockPromptTarget @@ -1771,6 +1772,7 @@ async def hold_async(*, conversation_id: str, **_: Any) -> None: release_first.set() await active + @pytest.mark.timeout(90) async def test_target_pacing_does_not_block_another_target_async( self, *, @@ -1785,7 +1787,9 @@ async def test_target_pacing_does_not_block_another_target_async( ) await sqlite_instance.add_attack_results_to_memory_async(attack_results=[other_attack]) waiting, release = asyncio.Event(), asyncio.Event() + metadata_started, release_metadata = asyncio.Event(), asyncio.Event() delays: list[float] = [] + update_attack_result = sqlite_instance.update_attack_result_by_id_async async def pace_async(delay: float) -> None: delays.append(delay) @@ -1795,11 +1799,18 @@ async def pace_async(delay: float) -> None: else: assert delay == 1 + async def update_attack_result_async(*, attack_result_id: str, update_fields: dict[str, Any]) -> bool: + if attack_result_id == other_attack.attack_result_id: + metadata_started.set() + await release_metadata.wait() + return await update_attack_result(attack_result_id=attack_result_id, update_fields=update_fields) + other_request = _request(conversation_id=other_attack.conversation_id) other_request.target_registry_name = "other" with ( patch("pyrit.backend.services.message_send_service.get_target_service") as registry, patch("pyrit.prompt_target.common.utils.asyncio.sleep", side_effect=pace_async), + patch.object(sqlite_instance, "update_attack_result_by_id_async", side_effect=update_attack_result_async), ): registry.return_value.get_target_object.side_effect = lambda *, target_registry_name: ( other_target if target_registry_name == "other" else target @@ -1809,19 +1820,31 @@ async def pace_async(delay: float) -> None: attack_result_id=ar.attack_result_id, request=_request(conversation_id=ar.conversation_id) ) ) + tasks = [active] try: - await waiting.wait() - await asyncio.wait_for( - service.add_message_async(attack_result_id=other_attack.attack_result_id, request=other_request), - timeout=3, + await asyncio.wait_for(waiting.wait(), timeout=30) + independent = asyncio.create_task( + service.add_message_async(attack_result_id=other_attack.attack_result_id, request=other_request) ) + tasks.append(independent) + await asyncio.wait_for(metadata_started.wait(), timeout=30) assert delays == [2.0, 1.0] assert target.prompt_sent == [] assert other_target.prompt_sent == ["Hello"] assert not active.done() + assert not independent.done() + release_metadata.set() + await wait_for_completion_async(future=independent) + assert not active.done() + release.set() + await wait_for_completion_async(future=active) finally: release.set() - await active + release_metadata.set() + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) assert target.prompt_sent == ["Hello"] @pytest.mark.parametrize("stage", ["request", "response"])