teng-lin--notebooklm-py/tests/unit/test_from_storage_dual_protocol.py

"""Tests for the dual-protocol wrapper returned by ``NotebookLMClient.from_storage``.

In v0.5.0 ``from_storage`` stopped being an async coroutine and started
returning ``_FromStorageContext`` — an awaitable async-context-manager
wrapper that supports BOTH:

- ``async with NotebookLMClient.from_storage(...) as client:`` (canonical,
  no ``DeprecationWarning``).
- ``async with await NotebookLMClient.from_storage(...) as client:`` /
  bare ``await NotebookLMClient.from_storage(...)`` (legacy, emits
  ``DeprecationWarning`` naming the v1.0 removal).

This module covers both protocols, the deprecation message text, the
type identity of the yielded value, and the lazy-load behavior that
``_FromStorageContext`` itself does no I/O until used.
"""

from __future__ import annotations

import json
import warnings
from pathlib import Path

import pytest
from pytest_httpx import HTTPXMock

from notebooklm.client import NotebookLMClient, _FromStorageContext


def _write_storage_state(tmp_path: Path) -> Path:
    """Drop a minimal storage_state.json on disk and return the path."""
    storage_file = tmp_path / "storage_state.json"
    storage_state = {
        "cookies": [
            {"name": "SID", "value": "dual_sid", "domain": ".google.com"},
            {
                "name": "__Secure-1PSIDTS",
                "value": "dual_1psidts",
                "domain": ".google.com",
            },
            {"name": "HSID", "value": "dual_hsid", "domain": ".google.com"},
        ],
        "origins": [],
    }
    storage_file.write_text(json.dumps(storage_state))
    return storage_file


def _stub_homepage(httpx_mock: HTTPXMock) -> None:
    """Stub the notebooklm homepage GET used to seed the CSRF/session tokens."""
    html = '"SNlM0e":"dual_csrf" "FdrFJe":"dual_session"'
    httpx_mock.add_response(
        url="https://notebooklm.google.com/",
        content=html.encode(),
    )


class TestCanonicalAsyncWith:
    """The new canonical path: ``async with from_storage(...) as client:``."""

    @pytest.mark.asyncio
    async def test_async_with_yields_connected_client(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """No ``await`` on the call; ``__aenter__`` builds + opens the client."""
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        async with NotebookLMClient.from_storage(path=str(storage_file)) as client:
            assert isinstance(client, NotebookLMClient)
            assert client.is_connected is True
            # The auth dropped onto the wrapper survives the lazy build.
            assert client.auth.cookies[("SID", ".google.com", "/")] == "dual_sid"

        # Exiting the context manager tears the session down again.
        assert client.is_connected is False

    @pytest.mark.asyncio
    async def test_async_with_emits_no_deprecation_warning(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """The canonical path is the migration target — must be warning-free."""
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            async with NotebookLMClient.from_storage(path=str(storage_file)) as client:
                _ = client  # nothing — just want a clean entry/exit
        deprecations = [w for w in caught if issubclass(w.category, DeprecationWarning)]
        assert deprecations == [], (
            f"Canonical `async with from_storage(...)` must not emit "
            f"DeprecationWarning; got: {[str(w.message) for w in deprecations]}"
        )

    @pytest.mark.asyncio
    async def test_async_with_forwards_chat_response_cap_override(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """Lazy from-storage build forwards chat response byte cap kwargs."""
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        async with NotebookLMClient.from_storage(
            path=str(storage_file),
            chat_response_max_bytes=123456,
        ) as client:
            assert client.chat._chat_response_max_bytes == 123456


class TestLegacyAwaitForm:
    """The legacy path: ``await from_storage(...)`` / ``async with await ...``."""

    @pytest.mark.asyncio
    async def test_await_still_returns_client(self, tmp_path: Path, httpx_mock: HTTPXMock) -> None:
        """Bare ``await`` returns a built-but-unentered client (legacy contract)."""
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        with pytest.warns(DeprecationWarning, match="removed in v1.0"):
            client = await NotebookLMClient.from_storage(path=str(storage_file))

        assert isinstance(client, NotebookLMClient)
        # ``from_storage`` historically returned an unentered client; the
        # legacy path preserves that contract so existing call sites that
        # then do ``async with client:`` keep working.
        assert client.is_connected is False

    @pytest.mark.asyncio
    async def test_async_with_await_pipes_through(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """``async with await from_storage(...)`` still works end-to-end."""
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        with pytest.warns(DeprecationWarning, match="removed in v1.0"):
            unentered = await NotebookLMClient.from_storage(path=str(storage_file))
        async with unentered as client:
            assert isinstance(client, NotebookLMClient)
            assert client.is_connected is True
        assert client.is_connected is False

    @pytest.mark.asyncio
    async def test_deprecation_warning_mentions_v1_removal(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """The warning must name v1.0 so users have a migration target."""
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            await NotebookLMClient.from_storage(path=str(storage_file))

        deprecations = [w for w in caught if issubclass(w.category, DeprecationWarning)]
        assert deprecations, "Expected at least one DeprecationWarning"
        msg = str(deprecations[0].message)
        assert "removed in v1.0" in msg, (
            f"DeprecationWarning must include 'removed in v1.0' for migration clarity; got: {msg!r}"
        )
        # The message should also point users at the new idiom.
        assert "async with" in msg, (
            f"DeprecationWarning should point at the canonical idiom; got: {msg!r}"
        )


class TestWrapperShape:
    """Properties of the wrapper itself, independent of either protocol."""

    def test_from_storage_is_sync(self, tmp_path: Path) -> None:
        """``from_storage`` no longer returns a coroutine — it returns a wrapper.

        Crucially, calling it does NOT await anything: it just constructs
        ``_FromStorageContext`` with the kwargs captured. No I/O happens
        until ``__aenter__`` or ``__await__``.
        """
        # No event loop needed; calling ``from_storage`` is sync.
        wrapper = NotebookLMClient.from_storage(
            path=str(tmp_path / "nonexistent.json"),
        )
        assert isinstance(wrapper, _FromStorageContext)
        # No coroutine, so nothing to await — and importantly, the missing
        # storage file has NOT been read yet (no FileNotFoundError).

    @pytest.mark.asyncio
    async def test_yielded_value_is_notebooklmclient(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """Type-check assertion (runtime): ``async with ... as client`` is a client.

        ``_FromStorageContext.__aenter__`` is typed to return
        ``NotebookLMClient`` so static type-checkers (mypy / pyright)
        narrow the ``as`` binding to ``NotebookLMClient``. We assert the
        runtime side here; mypy enforces the static side via the type
        annotation on ``__aenter__``.
        """
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        async with NotebookLMClient.from_storage(path=str(storage_file)) as client:
            assert isinstance(client, NotebookLMClient)

    @pytest.mark.asyncio
    async def test_build_is_idempotent_when_called_twice(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """Awaiting and then entering the same wrapper must not double-load auth.

        Pathological but legal sequence: ``client = await wrapper`` (legacy
        path returns the unentered client), then ``async with wrapper as
        client2:`` enters via the same wrapper. Both call sites should
        see the same ``NotebookLMClient`` instance, and the auth load
        should run exactly once.
        """
        storage_file = _write_storage_state(tmp_path)
        # Only one homepage hit allowed — proves auth load doesn't re-run.
        _stub_homepage(httpx_mock)

        wrapper = NotebookLMClient.from_storage(path=str(storage_file))

        with pytest.warns(DeprecationWarning, match="removed in v1.0"):
            client_from_await = await wrapper

        async with wrapper as client_from_aenter:
            assert client_from_aenter is client_from_await
            assert client_from_aenter.is_connected is True

    @pytest.mark.asyncio
    async def test_async_with_propagates_body_exceptions(
        self, tmp_path: Path, httpx_mock: HTTPXMock
    ) -> None:
        """Exceptions inside the ``async with`` body must propagate.

        ``_FromStorageContext.__aexit__`` returns ``None`` (falsy), so
        Python's async-with protocol re-raises the body exception. We
        also confirm the underlying client gets torn down so the test
        is asserting end-to-end teardown, not just the wrapper-level
        no-suppress contract.
        """
        storage_file = _write_storage_state(tmp_path)
        _stub_homepage(httpx_mock)

        sentinel: list[NotebookLMClient] = []
        with pytest.raises(ValueError, match="body raised"):
            async with NotebookLMClient.from_storage(path=str(storage_file)) as client:
                sentinel.append(client)
                assert client.is_connected is True
                raise ValueError("body raised")

        # Even though the body raised, ``__aexit__`` on the wrapper still
        # closed the underlying client.
        assert sentinel[0].is_connected is False