TEST: Stabilize target startup and analytics deadline tests

doc/contributing/4_running_tests.md

@@ -78,6 +78,12 @@ 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.
 
+The discovery provider-state test and Copilot cancellation-test startup use the
+shared `TEST_WATCHDOG_SECONDS` backstop in `tests/unit/async_utils.py`. Discovery
+timeout coverage expires the real cancellation path only after the request starts.
+Event-controlled analytics expiry tests freeze both deadline clocks so floating-point
+rounding cannot turn a valid query budget into a failed assertion.
+
 For scenario correctness tests, wait for the owned task to finish, including
 persistence and scheduler handoff. Use events to control concurrent steps and a
 finite test watchdog to detect deadlocks, not short per-step latency limits.

tests/unit/async_utils.py

@@ -6,13 +6,16 @@
 
 T = TypeVar("T")
 
+# Deadlock backstop for correctness tests, not a latency budget.
+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

tests/unit/backend/test_analytics_lifecycle.py

@@ -163,7 +163,7 @@ async def test_http_timeout_keeps_native_cleanup_busy_until_runtime_drain_async(
     cleanup = held_native_cleanup
 
     async def expire_async(tasks: set[asyncio.Task[Any]], *, timeout: float) -> tuple[set[Any], set[Any]]:
-        assert 0 < timeout <= 300
+        assert timeout == 300.0
         await expire.wait()
         return set(), tasks
 
@@ -177,12 +177,15 @@ async def close_async() -> None:
         transport=httpx.ASGITransport(app=analytics_runtime.app), base_url="http://test"
     ) as client:
         with (
+            patch("pyrit.analytics._execution.monotonic", return_value=100.0),
+            patch("pyrit.memory.query_control.time") as clock,
             patch("pyrit.analytics._execution.asyncio.wait", side_effect=expire_async),
             patch.object(analytics_service, "close_async", side_effect=close_async),
             patch.object(
                 sqlite_instance, "dispose_engine_async", wraps=sqlite_instance.dispose_engine_async
             ) as dispose,
         ):
+            clock.monotonic.return_value = 100.0
             request = asyncio.create_task(client.post("/api/analytics/attacks/results", json={}))
             shutdown: asyncio.Task[None] | None = None
             try:

tests/unit/prompt_target/target/test_github_copilot_target.py

@@ -15,7 +15,7 @@
 from uuid import UUID, uuid4
 
 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, wait_for_completion_async
 from unit.mocks import store_message_async
 
 from pyrit.models import Message, MessagePiece, MessageScorable, ScoringExpectation
@@ -1203,7 +1203,7 @@ async def send_and_wait_async(*_args: Any, **_kwargs: Any) -> None:
     )
 
     try:
-        await asyncio.wait_for(send_started.wait(), timeout=2.0)
+        await asyncio.wait_for(send_started.wait(), timeout=TEST_WATCHDOG_SECONDS)
         send_task.cancel()
         with pytest.raises(asyncio.CancelledError) as exc_info:
             await asyncio.wait_for(send_task, timeout=2.0)

tests/unit/prompt_target/test_discover_target_capabilities.py

@@ -4,6 +4,7 @@
 import asyncio
 import json
 import logging
+from collections.abc import Awaitable
 from copy import deepcopy
 from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -12,6 +13,7 @@
 import pytest
 from mcp.types import Tool as MCPToolDefinition
 from openai.types.chat import ChatCompletion
+from unit.async_utils import TEST_WATCHDOG_SECONDS
 from unit.mocks import MockPromptTarget
 
 from pyrit.models import Message, MessagePiece, PromptDataType, RequestTraceContext
@@ -792,15 +794,34 @@ async def direct_async() -> str:
         original_value_data_type="function_call",
     ).to_message()
 
+    requested = asyncio.Event()
+    cancelled = asyncio.Event()
+
     async def finish_async(**kwargs: object) -> Message:
         if outcome == "error":
             raise ValueError("Probe request failed")
         if outcome == "cancel":
             raise asyncio.CancelledError
         if outcome == "timeout":
-            await asyncio.Event().wait()
+            requested.set()
+            try:
+                await asyncio.Event().wait()
+            finally:
+                cancelled.set()
         return response
 
+    original_wait_for = asyncio.wait_for
+
+    async def expire_once_requested_async(awaitable: Awaitable[list[Message]], *, timeout: float) -> list[Message]:
+        probe = asyncio.ensure_future(awaitable)
+        try:
+            await original_wait_for(requested.wait(), timeout=TEST_WATCHDOG_SECONDS)
+            return await original_wait_for(probe, timeout=0)
+        finally:
+            if not probe.done():
+                probe.cancel()
+            await asyncio.gather(probe, return_exceptions=True)
+
     with (
         patch.object(provider, "execution_scope_async") as scope,
         patch.object(provider, "get_tools_async", new_callable=AsyncMock) as discover,
@@ -816,13 +837,17 @@ async def probe_async() -> None:
                 capabilities=[CapabilityName.JSON_OUTPUT],
                 test_modalities=set(),
                 retries=0,
-                # The deadline includes async memory access before the mocked request.
-                per_probe_timeout_s=5.0,
+                per_probe_timeout_s=TEST_WATCHDOG_SECONDS,
             )
 
         if outcome == "cancel":
             with pytest.raises(asyncio.CancelledError):
                 await probe_async()
+        elif outcome == "timeout":
+            with patch.object(asyncio, "wait_for", side_effect=expire_once_requested_async) as deadline:
+                await probe_async()
+            deadline.assert_awaited_once()
+            assert cancelled.is_set()
         else:
             await probe_async()
         send.assert_awaited_once()