test: make timing-sensitive budget and task-timeout tests deterministic

tests/test_permission_denied_budget.py

@@ -19,6 +19,7 @@
 import asyncio
 import logging
 import time
+from types import SimpleNamespace
 
 import pytest
 
@@ -28,6 +29,7 @@
     BuyerAgentRegistry,
     HttpSigCredential,
 )
+from adcp.decisioning import _permission_denied_budget as budget_module
 from adcp.decisioning._permission_denied_budget import (
     BUDGET_ENV_VAR,
     DEFAULT_BUDGET_MS,
@@ -91,6 +93,39 @@ def test_env_var_invalid_still_falls_back(monkeypatch, caplog):
 # ---------------------------------------------------------------------------
 
 
+class _VirtualClock:
+    """Deterministic time for the budget module.
+
+    Upper bounds and differences between real elapsed times depend on how
+    busy the machine is (``pytest -n auto`` saturates it), so the tests that
+    assert them run the budget on this clock instead: ``perf_counter`` reads
+    it, and ``sleep`` and the fake registry advance it, exactly.
+    """
+
+    def __init__(self) -> None:
+        self.now = 1000.0
+        self.sleeps: list[float] = []
+
+    def perf_counter(self) -> float:
+        return self.now
+
+    async def sleep(self, seconds: float) -> None:
+        self.sleeps.append(seconds)
+        self.now += max(seconds, 0.0)
+        await asyncio.sleep(0)
+
+    def elapsed_ms_since(self, start: float) -> float:
+        return (self.now - start) * 1000.0
+
+
+@pytest.fixture
+def virtual_clock(monkeypatch) -> _VirtualClock:
+    clock = _VirtualClock()
+    monkeypatch.setattr(budget_module, "time", SimpleNamespace(perf_counter=clock.perf_counter))
+    monkeypatch.setattr(budget_module, "asyncio", SimpleNamespace(sleep=clock.sleep))
+    return clock
+
+
 @pytest.mark.asyncio
 async def test_enforce_sleeps_to_budget_when_branch_is_fast(monkeypatch):
     """A branch that does no work waits the full budget."""
@@ -105,16 +140,14 @@ async def test_enforce_sleeps_to_budget_when_branch_is_fast(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_enforce_does_not_sleep_when_already_overrun(monkeypatch):
+async def test_enforce_does_not_sleep_when_already_overrun(monkeypatch, virtual_clock):
     """A branch that already exceeded the budget passes through immediately."""
     monkeypatch.setenv(BUDGET_ENV_VAR, "5")
     budget = PermissionDeniedBudget()
     # Burn more than the budget before enforce().
-    await asyncio.sleep(0.020)
-    start = time.perf_counter()
+    virtual_clock.now += 0.020
     await budget.enforce()
-    elapsed_ms = (time.perf_counter() - start) * 1000.0
-    assert elapsed_ms < 5.0, f"enforce() should be a no-op once overrun, got {elapsed_ms:.1f} ms"
+    assert virtual_clock.sleeps == []
 
 
 @pytest.mark.asyncio
@@ -146,18 +179,30 @@ async def test_inf_does_not_hang(monkeypatch):
 
 
 class _SlowRegistry(BuyerAgentRegistry):
-    """A registry whose resolve methods take a configurable amount of time."""
+    """A registry whose resolve methods take a configurable amount of time.
+
+    With ``clock``, the delay advances that virtual clock instead of sleeping.
+    """
 
-    def __init__(self, delay_ms: float, agent: BuyerAgent | None) -> None:
+    def __init__(
+        self, delay_ms: float, agent: BuyerAgent | None, *, clock: _VirtualClock | None = None
+    ) -> None:
         self._delay_seconds = delay_ms / 1000.0
         self._agent = agent
+        self._clock = clock
+
+    async def _delay(self) -> None:
+        if self._clock is not None:
+            await self._clock.sleep(self._delay_seconds)
+        else:
+            await asyncio.sleep(self._delay_seconds)
 
     async def resolve_by_agent_url(self, agent_url: str) -> BuyerAgent | None:
-        await asyncio.sleep(self._delay_seconds)
+        await self._delay()
         return self._agent
 
     async def resolve_by_credential(self, credential):  # pragma: no cover - not used here
-        await asyncio.sleep(self._delay_seconds)
+        await self._delay()
         return self._agent
 
 
@@ -245,42 +290,32 @@ async def test_active_branch_skips_budget(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_branch_parity_absorbs_registry_variance(monkeypatch):
+async def test_branch_parity_absorbs_registry_variance(monkeypatch, virtual_clock):
     """The whole point: fast registry-miss + slow real-row read should
-    both come out at ≈ budget, not (0 ms, registry_delay)."""
+    both come out at the budget, not (0 ms, registry_delay)."""
     monkeypatch.setenv(BUDGET_ENV_VAR, "50")
 
-    # Branch 1: registry-miss path with a fast cache-miss-returning-None.
-    fast_miss = _SlowRegistry(delay_ms=1.0, agent=None)
-    miss_samples: list[float] = []
-    for _ in range(5):
-        start = time.perf_counter()
-        with pytest.raises(AdcpError):
-            await _resolve_buyer_agent(fast_miss, _http_sig_auth("https://buyer.example"))
-        miss_samples.append((time.perf_counter() - start) * 1000.0)
+    async def elapsed_ms(registry: _SlowRegistry) -> float:
+        start = virtual_clock.now
+        with pytest.raises(AdcpError) as exc:
+            await _resolve_buyer_agent(registry, _http_sig_auth("https://buyer.example"))
+        assert exc.value.code == "PERMISSION_DENIED"
+        return virtual_clock.elapsed_ms_since(start)
 
+    # Branch 1: registry-miss path with a fast cache-miss-returning-None.
+    fast_miss = _SlowRegistry(delay_ms=1.0, agent=None, clock=virtual_clock)
     # Branch 2: unknown-status path with a slower registry read for a real row.
     unknown_agent = BuyerAgent(
         agent_url="https://buyer.example",
         display_name="Test Buyer",
         status="pending_review",
     )
-    slow_hit = _SlowRegistry(delay_ms=10.0, agent=unknown_agent)
-    hit_samples: list[float] = []
-    for _ in range(5):
-        start = time.perf_counter()
-        with pytest.raises(AdcpError):
-            await _resolve_buyer_agent(slow_hit, _http_sig_auth("https://buyer.example"))
-        hit_samples.append((time.perf_counter() - start) * 1000.0)
-
-    miss_max = max(miss_samples)
-    hit_max = max(hit_samples)
-    # Both should be dominated by the 50 ms budget. The branch difference
-    # must be much smaller than the registry-delay difference (9 ms) that
-    # would leak without the budget.
-    branch_delta = abs(hit_max - miss_max)
-    assert (
-        branch_delta < 15.0
-    ), f"branch parity broken: hit={hit_max:.1f} ms vs miss={miss_max:.1f} ms"
-    # Both should be at or above the budget.
-    assert min(miss_samples + hit_samples) >= 45.0
+    slow_hit = _SlowRegistry(delay_ms=10.0, agent=unknown_agent, clock=virtual_clock)
+
+    # Both branches end exactly at the 50 ms budget: the 9 ms registry-delay
+    # difference that would leak without the budget is absorbed into it.
+    assert await elapsed_ms(fast_miss) == pytest.approx(50.0)
+    assert await elapsed_ms(slow_hit) == pytest.approx(50.0)
+    # A registry read slower than the budget is not padded further.
+    slower_than_budget = _SlowRegistry(delay_ms=80.0, agent=unknown_agent, clock=virtual_clock)
+    assert await elapsed_ms(slower_than_budget) == pytest.approx(80.0)

tests/test_task_options.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import inspect
 import time
 import warnings
@@ -14,11 +15,10 @@
 from pydantic import BaseModel, TypeAdapter
 
 from adcp import ADCPClient, TaskOptions, TaskRecoveryMetadata, _idempotency
-from adcp._null_clear import nullable_request_paths
 from adcp.exceptions import ADCPTimeoutError
 from adcp.protocols.base import ProtocolAdapter
 from adcp.protocols.mcp import MCPAdapter
-from adcp.task_options import mark_task_dispatched
+from adcp.task_options import _TaskDeadlineExpiredError, mark_task_dispatched
 from adcp.types import GetProductsRequest
 from adcp.types.core import ActivityType, AgentConfig, Protocol, TaskResult, TaskStatus
 
@@ -138,11 +138,48 @@ async def delayed(_params: dict[str, Any]) -> TaskResult[Any]:
             await task
 
 
+# Generous enough that pre-dispatch work (cold schema loads, validation) can
+# never exhaust it, however loaded the machine is.
+_PAST_ANY_PREFLIGHT = 60.0
+
+
+@contextlib.contextmanager
+def _deadline_expires_on(dispatched: asyncio.Event):
+    """Expire the task deadline exactly when ``dispatched`` is set.
+
+    A real 10-100 ms deadline also covers the work before dispatch, which is
+    slow when cold or under ``pytest -n auto``: the deadline then expires
+    before dispatch and there is nothing to recover. Replacing the timer seam
+    (``run_with_timeout``) makes the expiry land after dispatch every time.
+    The client's deadline state, dispatch marking and recovery are all real,
+    and the expiry raises what the real timer raises; other tests here
+    exercise the real timer.
+    """
+
+    async def expire_on_dispatch(factory: Any, _timeout: float) -> Any:
+        work = asyncio.ensure_future(factory())
+        signal = asyncio.ensure_future(dispatched.wait())
+        try:
+            await asyncio.wait({work, signal}, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            signal.cancel()
+        if work.done():
+            return work.result()
+        work.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await work
+        raise _TaskDeadlineExpiredError
+
+    with patch("adcp.client.run_with_timeout", new=expire_on_dispatch):
+        yield
+
+
 @pytest.mark.asyncio
 async def test_mutation_timeout_carries_secret_safe_recovery() -> None:
     activities = []
     client = _client(on_activity=activities.append)
     key = "0123456789abcdef-recovery-key"
+    dispatched_event = asyncio.Event()
 
     async def dispatched(_params: dict[str, Any]) -> TaskResult[Any]:
         mark_task_dispatched(
@@ -151,18 +188,20 @@ async def dispatched(_params: dict[str, Any]) -> TaskResult[Any]:
             mutating=True,
             idempotency_key=key,
         )
-        await asyncio.sleep(1)
+        dispatched_event.set()
+        await asyncio.Event().wait()
         return _completed()
 
     with (
         warnings.catch_warnings(),
         patch.object(client.adapter, "build_creative", new=dispatched),
+        _deadline_expires_on(dispatched_event),
     ):
         warnings.simplefilter("ignore", DeprecationWarning)
         with pytest.raises(ADCPTimeoutError) as caught:
             await client.build_creative_legacy(
                 _Request(idempotency_key=key),
-                options=TaskOptions(timeout=0.01),  # type: ignore[arg-type]
+                options=TaskOptions(timeout=_PAST_ANY_PREFLIGHT),  # type: ignore[arg-type]
             )
 
     error = caught.value
@@ -182,10 +221,8 @@ async def test_new_mutations_timeout_with_recovery_metadata(task_name: str) -> N
     client = _client()
     key = f"0123456789abcdef-{task_name}"
     dispatched_event = asyncio.Event()
-    # This case exercises a timeout AFTER dispatch. Cold nullable-schema loading
-    # can itself exhaust the budget before dispatch, correctly leaving recovery
-    # unset; the pre-dispatch deadline case below covers that separate contract.
-    nullable_request_paths(task_name, client.get_adcp_version())
+    # This case exercises a timeout AFTER dispatch; the pre-dispatch deadline
+    # case below covers that separate contract.
 
     async def dispatched(params: dict[str, Any]) -> TaskResult[Any]:
         mark_task_dispatched(
@@ -198,11 +235,14 @@ async def dispatched(params: dict[str, Any]) -> TaskResult[Any]:
         await asyncio.Event().wait()
         return _completed()
 
-    with patch.object(client.adapter, task_name, new=dispatched):
+    with (
+        patch.object(client.adapter, task_name, new=dispatched),
+        _deadline_expires_on(dispatched_event),
+    ):
         with pytest.raises(ADCPTimeoutError) as caught:
             await getattr(client, task_name)(
                 _Request(idempotency_key=key),
-                options=TaskOptions(timeout=0.1),
+                options=TaskOptions(timeout=_PAST_ANY_PREFLIGHT),
             )
 
     assert dispatched_event.is_set()
@@ -215,6 +255,8 @@ async def dispatched(params: dict[str, Any]) -> TaskResult[Any]:
 async def test_real_mcp_dispatch_reports_the_exact_pinned_wire_key() -> None:
     from adcp.server.responses import capabilities_response
 
+    on_the_wire = asyncio.Event()
+
     class BlockingSession:
         sent: dict[str, Any] | None = None
 
@@ -224,19 +266,24 @@ async def call_tool(self, name: str, params: dict[str, Any]) -> Any:
                 caps["adcp"]["idempotency"] = {"supported": False}
                 return SimpleNamespace(isError=False, structuredContent=caps, content=[])
             self.sent = params
-            await asyncio.sleep(1)
+            on_the_wire.set()
+            await asyncio.Event().wait()
 
     session = BlockingSession()
     client = ADCPClient.from_mcp_client(session, agent_id="in-process")  # type: ignore[arg-type]
     client.adapter.request_validation_mode = "off"
     await client.fetch_capabilities()
     key = "0123456789abcdef-pinned-wire-key"
 
-    with warnings.catch_warnings(), client.use_idempotency_key(key):
+    with (
+        warnings.catch_warnings(),
+        client.use_idempotency_key(key),
+        _deadline_expires_on(on_the_wire),
+    ):
         warnings.simplefilter("ignore", DeprecationWarning)
         with pytest.raises(ADCPTimeoutError) as caught:
             await client.build_creative_legacy(  # type: ignore[arg-type]
-                _NoKeyRequest(), options=TaskOptions(timeout=0.05)
+                _NoKeyRequest(), options=TaskOptions(timeout=_PAST_ANY_PREFLIGHT)
             )
 
     assert session.sent is not None