diff --git a/tests/test_permission_denied_budget.py b/tests/test_permission_denied_budget.py index 0da826e8f..3b810567e 100644 --- a/tests/test_permission_denied_budget.py +++ b/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) diff --git a/tests/test_task_options.py b/tests/test_task_options.py index 790074f3c..4378c28e2 100644 --- a/tests/test_task_options.py +++ b/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,7 +266,8 @@ 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] @@ -232,11 +275,15 @@ async def call_tool(self, name: str, params: dict[str, Any]) -> Any: 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