Files
teng-lin--notebooklm-py/tests/integration/concurrency/conftest.py
T
wehub-resource-sync 09e9f3545f
Test / Code Quality (push) Has been cancelled
Test / Test (macos-latest, Python 3.10) (push) Has been cancelled
Test / Test (macos-latest, Python 3.11) (push) Has been cancelled
Test / Test (macos-latest, Python 3.12) (push) Has been cancelled
Test / Test (macos-latest, Python 3.13) (push) Has been cancelled
Test / Test (macos-latest, Python 3.14) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.10) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.11) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.12) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.13) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.14) (push) Has been cancelled
Test / Test (windows-latest, Python 3.10) (push) Has been cancelled
Test / Test (windows-latest, Python 3.11) (push) Has been cancelled
Test / Test (windows-latest, Python 3.12) (push) Has been cancelled
Test / Test (windows-latest, Python 3.13) (push) Has been cancelled
Test / Test (windows-latest, Python 3.14) (push) Has been cancelled
CodeQL / Analyze (push) Has been cancelled
dependency-audit / pip-audit (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:30:13 +08:00

397 lines
14 KiB
Python

"""Fixtures for the concurrency integration harness.
Three fixtures, all explicitly named and class-based — no clever hooks,
no monkeypatching. Fix PRs consume these to TDD red-green against
specific concurrency bugs (refresh races, semaphore gating,
cancellation propagation, etc.).
Fixtures
--------
``mock_transport_concurrent``
A class-based ``httpx.AsyncBaseTransport`` that records peak in-flight
concurrency and supports per-request controllable response timing.
See ``ConcurrentMockTransport`` for the full method surface.
``barrier_factory``
Returns a callable ``make_barrier(n)`` that produces a
``EventBarrier`` for deterministic interleaving of N coroutines.
Each barrier is a one-shot synchronization point: every arriver
awaits ``arrive()``; once N have arrived, all are released.
Built on ``asyncio.Event`` (not ``asyncio.Barrier``) so behavior is
identical on Python 3.10 (the project's minimum supported version).
``cancellation_helper``
Wraps a coroutine in ``asyncio.wait_for`` and emits a structured
diagnostic on cancellation: which coroutine label, what timeout,
and the captured traceback. Used by per-fix tests that need to
distinguish "the bug under test deadlocked" from "the test timed
out for an unrelated reason."
Non-goals
---------
- NOT a load tester. Peak-inflight assertions are coarse (>= 80 of 100)
because asyncio task scheduling is not perfectly parallel.
- NOT a property-based generator (no Hypothesis dep added).
- NOT a thread-pool stress harness (no real threads — pure asyncio).
- pytest-xdist + asyncio: each xdist worker has its own event loop, so
fixture state is never shared across workers. Tests that assert on
process-global state must mark themselves ``@pytest.mark.xdist_group``.
"""
from __future__ import annotations
import asyncio
import json
import logging
import traceback
from collections import deque
from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any, TypeVar
import httpx
import pytest
from notebooklm.rpc import RPCMethod
logger = logging.getLogger(__name__)
T = TypeVar("T")
# ---------------------------------------------------------------------------
# Default response payload
# ---------------------------------------------------------------------------
# A minimal valid batchexecute response that decodes to ``[]`` for
# ``LIST_NOTEBOOKS``. Reused by ``ConcurrentMockTransport`` whenever the
# response queue is empty so tests don't have to enqueue 100 identical
# responses for a 100-way fan-out.
_DEFAULT_RPC_ID = RPCMethod.LIST_NOTEBOOKS.value
def _default_rpc_response_text(rpc_id: str = _DEFAULT_RPC_ID) -> str:
"""Build a minimal valid batchexecute response that decodes to ``[]``."""
inner = json.dumps([])
chunk = json.dumps([["wrb.fr", rpc_id, inner, None, None]])
return f")]}}'\n{len(chunk)}\n{chunk}\n"
def install_post_as_stream(
monkeypatch: pytest.MonkeyPatch | None,
http_client: Any,
fake_post: Callable[..., Awaitable[Any]],
) -> None:
"""Adapt legacy fake ``post`` callbacks to the streaming RPC POST API."""
@asynccontextmanager
async def fake_stream(method: str, url: str, **kwargs: Any) -> Any:
response = await fake_post(url, **kwargs)
if type(response) is httpx.Response:
yield response
return
text = getattr(response, "text", "")
payload = text.encode("utf-8") if isinstance(text, str) else bytes(text or b"")
raw_status = getattr(response, "status_code", 200)
status = raw_status if isinstance(raw_status, int) else 200
try:
raw_headers = getattr(response, "headers", None)
except AttributeError:
raw_headers = None
try:
headers = dict(raw_headers) if raw_headers else None
except (TypeError, AttributeError):
headers = None
yield httpx.Response(
status_code=status,
headers=headers,
content=payload,
request=httpx.Request("POST", url),
)
if monkeypatch is not None:
monkeypatch.setattr(http_client, "stream", fake_stream)
else:
http_client.stream = fake_stream
# ---------------------------------------------------------------------------
# ConcurrentMockTransport
# ---------------------------------------------------------------------------
@dataclass
class _InflightTracker:
"""Plain counter pair. Single-threaded asyncio — no lock needed."""
current: int = 0
peak: int = 0
def enter(self) -> None:
self.current += 1
if self.current > self.peak:
self.peak = self.current
def exit(self) -> None:
# Hard assert: a double-`finally` or transport-reuse bug would
# silently drive `current` negative and quietly invalidate every
# peak assertion downstream. As a test helper we want loud
# failure, not silent corruption.
assert self.current > 0, "exit() called more times than enter()"
self.current -= 1
class ConcurrentMockTransport(httpx.AsyncBaseTransport):
"""Mock transport that records concurrent in-flight requests.
Designed for asyncio fan-out tests: every ``handle_async_request``
increments an in-flight counter, awaits a configurable delay (so
sibling tasks can pile up at the same await point), then decrements
and returns a queued (or default) response.
Methods
-------
queue_response(response_or_factory)
Append a response to the FIFO queue. May be:
- ``httpx.Response`` instance.
- Tuple ``(status_code, text)`` for convenience.
- Callable ``(httpx.Request) -> httpx.Response`` for per-request
shaping (e.g. echoing the URL, returning an error for one URL).
When the queue is empty, ``_default_rpc_response_text`` is used.
set_delay(seconds)
Set the per-request artificial delay (default ``0.05``s — long
enough that a 100-way ``asyncio.gather`` reliably stacks all 100
callers at the await point before any returns).
get_inflight_count()
Current in-flight request count.
get_peak_inflight()
High-water mark observed since construction (or last ``reset()``).
request_count()
Total requests served so far.
captured_requests()
Snapshot of every request observed (for assertions on URL,
headers, body — useful in per-fix PRs).
reset()
Clear counters, queued responses, and captured request history.
The configured ``_delay`` is preserved (a session-scoped consumer
wants its delay choice to survive reset). Call ``set_delay(...)``
explicitly if you need to reset the timing too.
Thread-safety note
------------------
All state mutation happens on the asyncio event loop's single thread.
No locks are needed. If a future test wires this into a real
threadpool, add an ``asyncio.Lock`` around the counter mutations.
"""
# Type alias kept inline so readers don't have to scroll for it.
_ResponseFactory = Callable[[httpx.Request], httpx.Response]
_QueuedResponse = httpx.Response | tuple[int, str] | _ResponseFactory
def __init__(self, *, default_delay: float = 0.05) -> None:
self._tracker = _InflightTracker()
# ``deque`` for O(1) FIFO popleft. ``list.pop(0)`` shifts the whole
# backing array each dequeue; immaterial at ~100 items but the wrong
# data structure for a queue.
self._queue: deque[ConcurrentMockTransport._QueuedResponse] = deque()
self._delay: float = default_delay
self._captured: list[httpx.Request] = []
self._request_count: int = 0
# -- configuration -------------------------------------------------
def queue_response(self, response_or_factory: _QueuedResponse) -> None:
"""Append a response to the FIFO queue.
Acceptable shapes documented on the class docstring.
"""
self._queue.append(response_or_factory)
def set_delay(self, seconds: float) -> None:
"""Set the artificial per-request delay.
``0`` is allowed but defeats the purpose of fan-out tests:
without a delay, requests complete before the next coroutine
even enters the transport. Use a small positive value (50ms is
the default) for fan-out work.
"""
if seconds < 0:
raise ValueError(f"delay must be >= 0, got {seconds}")
self._delay = seconds
# -- observation ---------------------------------------------------
def get_inflight_count(self) -> int:
return self._tracker.current
def get_peak_inflight(self) -> int:
return self._tracker.peak
def request_count(self) -> int:
return self._request_count
def captured_requests(self) -> list[httpx.Request]:
# Defensive copy so callers can iterate without racing future
# in-flight requests (which would mutate the underlying list
# via append).
return list(self._captured)
def reset(self) -> None:
self._tracker = _InflightTracker()
self._queue.clear()
self._captured.clear()
self._request_count = 0
# -- AsyncBaseTransport ABI ---------------------------------------
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
self._captured.append(request)
self._request_count += 1
self._tracker.enter()
try:
if self._delay > 0:
# The yield point that lets sibling coroutines stack up
# in the in-flight counter. Every concurrent caller hits
# ``enter()`` synchronously before any reaches this
# ``await``, so peak-inflight reflects the gather width.
await asyncio.sleep(self._delay)
return self._next_response(request)
finally:
self._tracker.exit()
# -- internal ------------------------------------------------------
def _next_response(self, request: httpx.Request) -> httpx.Response:
if not self._queue:
return httpx.Response(200, text=_default_rpc_response_text())
item = self._queue.popleft()
if isinstance(item, httpx.Response):
return item
if isinstance(item, tuple):
status, text = item
return httpx.Response(status, text=text)
if callable(item):
return item(request)
raise TypeError(
f"Unsupported queued response type: {type(item).__name__}. "
"Pass an httpx.Response, a (status, text) tuple, or a callable."
)
@pytest.fixture
def mock_transport_concurrent() -> ConcurrentMockTransport:
"""A fresh ``ConcurrentMockTransport`` per test."""
return ConcurrentMockTransport()
# ---------------------------------------------------------------------------
# barrier_factory
# ---------------------------------------------------------------------------
@dataclass
class EventBarrier:
"""One-shot N-arrival barrier built on ``asyncio.Event``.
Every arriver calls ``await arrive()``. The first ``N - 1`` arrivers
suspend at ``event.wait()``; the Nth arrival sets the event,
releasing all of them simultaneously.
Re-arming is intentionally NOT supported. Per-test barriers are
cheap; spawn a fresh one for each synchronization point. This
matches the guidance in the existing ``test_concurrency_refresh_race``
suite which uses one ``asyncio.Event`` per checkpoint.
"""
n: int
_event: asyncio.Event = field(default_factory=asyncio.Event)
_arrived: int = 0
async def arrive(self) -> None:
self._arrived += 1
if self._arrived >= self.n:
self._event.set()
await self._event.wait()
@property
def is_set(self) -> bool:
return self._event.is_set()
@property
def arrived_count(self) -> int:
return self._arrived
@pytest.fixture
def barrier_factory() -> Callable[[int], EventBarrier]:
"""Return a callable that builds N-arrival ``EventBarrier`` instances.
Usage::
async def test_thing(barrier_factory):
barrier = barrier_factory(3)
await asyncio.gather(
worker(barrier),
worker(barrier),
worker(barrier),
)
Implementation note: the factory is the fixture (not a barrier
instance) so a single test can spawn multiple independent
synchronization points without re-fixturing.
"""
def _make(n: int) -> EventBarrier:
if n <= 0:
raise ValueError(f"barrier arrivals must be >= 1, got {n}")
return EventBarrier(n=n)
return _make
# ---------------------------------------------------------------------------
# cancellation_helper
# ---------------------------------------------------------------------------
CancellationHelper = Callable[..., Awaitable[Any]]
@pytest.fixture
def cancellation_helper() -> CancellationHelper:
"""Wrap a coroutine in ``asyncio.wait_for`` with diagnostic on cancel.
Signature::
await cancellation_helper(coro, timeout=5.0, label="my-coro")
On ``asyncio.TimeoutError`` (Python 3.10) / ``TimeoutError`` (3.11+),
logs the label and timeout via the harness logger then re-raises.
The re-raise preserves test failure semantics — the helper's job is
to surface *which* coroutine deadlocked, not to swallow the failure.
The diagnostic is stderr-friendly and includes the traceback so a
CI failure with several pending coroutines is debuggable from the
log alone.
"""
async def _run(
coro: Awaitable[T],
*,
timeout: float = 5.0,
label: str = "<unlabeled>",
) -> T:
try:
return await asyncio.wait_for(coro, timeout=timeout)
except (TimeoutError, asyncio.TimeoutError):
tb = traceback.format_exc()
logger.error(
"cancellation_helper: coroutine %r timed out after %.3fs\n%s",
label,
timeout,
tb,
)
raise
return _run