"""Shared fixtures and helpers for tests/unit/. The ``make_core`` async context manager is imported directly by sibling test modules (e.g. ``from tests.unit.conftest import make_core``) now that the ``tests`` package chain is complete and fully qualified. """ from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager from typing import Any import httpx import pytest from notebooklm.auth import AuthTokens from tests._fixtures.kernel_test_helpers import install_http_client_for_test from tests._helpers.client_factory import build_client_shell_for_tests def install_post_as_stream( monkeypatch: pytest.MonkeyPatch | None, http_client: Any, fake_post: Callable[..., Awaitable[Any]], ) -> None: """Adapt a ``fake_post(...) -> Response`` mock to the streaming API. The RPC POST path uses :meth:`httpx.AsyncClient.stream` (so a running size guard can enforce :data:`notebooklm._streaming_post.MAX_RPC_RESPONSE_BYTES`). The bulk of the unit suite predates that switch and still expresses test intent as ``monkeypatch.setattr(client, "post", fake_post)``. This helper bridges the gap: it installs an ``async with client.stream(...)``-compatible fake on ``http_client.stream`` that delegates to the caller's existing ``fake_post`` (preserving call-count side effects, raised exceptions, and returned responses). For ``MagicMock`` responses lacking real ``aiter_bytes`` plumbing, the helper attaches a single-chunk async iterator over the mock's ``.text`` so the streaming wrapper's size-guard read loop terminates immediately. Real :class:`httpx.Response` instances are passed through unchanged — their built-in ``aiter_bytes`` works on already-buffered bodies. """ @asynccontextmanager async def fake_stream(method: str, url: str, **kwargs: Any) -> Any: # ``fake_post`` historically takes ``(url, **kwargs)`` — match that # call site exactly so existing argument-introspection in tests keeps # working unchanged. response = await fake_post(url, **kwargs) # ``type(...) is`` — not ``isinstance(...)`` — because ``MagicMock( # spec=httpx.Response)`` passes the isinstance check, which would # leave the streaming wrapper trying to read ``response.headers`` and # other spec-enforced attributes the test never set, raising # ``AttributeError`` deep inside production code instead of going # through the friendly rewrap branch below. if type(response) is httpx.Response: yield response return # Non-``httpx.Response`` (MagicMock-style): re-wrap into a real # :class:`httpx.Response` carrying the canned text so the streaming # wrapper's ``aiter_bytes`` + rebuild path works on it. Returning a # ``spec=httpx.Response`` MagicMock directly is brittle because the # spec rejects ad-hoc attribute access (``response.headers = {}`` # raises AttributeError), and the streaming wrapper reads several # attributes (``status_code``, ``headers``, ``request``, # ``aiter_bytes``) the mocks rarely set. 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) # MagicMock auto-mocks attributes, so ``status_code`` might be a Mock # whose ``__int__`` returns 1. Only treat real ints as set; otherwise # default to 200 (the canonical "no error" status the success-path # tests are implicitly asserting against). status = raw_status if isinstance(raw_status, int) else 200 # Preserve mock-set headers (e.g. ``retry-after``) so 429-path # ``exc.response.headers.get(...)`` introspection still returns the # value the test pinned. ``MagicMock(spec=...)`` raises AttributeError # for spec-defined names the test didn't explicitly set, so catch # that too — not just the missing-default case ``getattr`` already # handles. 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 wrapped = httpx.Response( status_code=status, headers=headers, content=payload, request=httpx.Request("POST", url), ) yield wrapped if monkeypatch is not None: monkeypatch.setattr(http_client, "stream", fake_stream) else: # MagicMock-style tests assign attributes directly rather than going # through ``monkeypatch.setattr``; honor that pattern for them. http_client.stream = fake_stream @pytest.fixture def auth_tokens(): """Create test authentication tokens for unit tests. Overrides the root-level fixture (single-cookie) with the full required cookie set so httpx_mock-based tests previously living in ``tests/integration/`` (later moved to ``tests/unit/``) can keep asserting on per-cookie wire values (e.g. ``SID=test_sid``, ``HSID=test_hsid``) without modification. The root fixture remains the canonical minimal jar for tests that don't inspect cookie headers. """ return AuthTokens( cookies={ "SID": "test_sid", "HSID": "test_hsid", "SSID": "test_ssid", "APISID": "test_apisid", "SAPISID": "test_sapisid", }, csrf_token="test_csrf_token", session_id="test_session_id", ) @asynccontextmanager async def make_core(refresh_callback=None, transport=None, refresh_retry_delay=0.0): """Yield an opened Session with optional mock transport; close cleanly. Args: refresh_callback: async callable returning ``AuthTokens`` (or raising) for use by ``_try_refresh_and_retry``. ``None`` skips refresh setup. transport: optional ``httpx.MockTransport`` so tests can observe the real ``httpx.Request`` after cookie merge. refresh_retry_delay: shortened in tests (default 0.0) to keep the suite fast — production default is 0.2s. """ auth = AuthTokens( csrf_token="CSRF_OLD", session_id="SID_OLD", cookies={"SID": "old_sid_cookie"}, ) core = build_client_shell_for_tests( auth=auth, refresh_callback=refresh_callback, refresh_retry_delay=refresh_retry_delay, ) await core.__aenter__() if transport is not None: # Replace the auto-built client with one that uses our transport so we # can observe real httpx.Request construction (cookie merge, headers). # Capture the cookie jar BEFORE aclose() — reading attributes off a # closed AsyncClient is brittle across httpx versions. prior_cookies = core._collaborators.kernel.get_http_client().cookies await core._collaborators.kernel.get_http_client().aclose() install_http_client_for_test( core._collaborators.kernel, httpx.AsyncClient( cookies=prior_cookies, transport=transport, timeout=httpx.Timeout(connect=1.0, read=5.0, write=5.0, pool=1.0), ), ) try: yield core finally: await core.close()