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

"""Regression guard for the ``CookiePersistence.save_lock`` contract.

Contract (documented at ``_cookie_persistence.py`` next to the lock definition):
``save_lock`` is acquired ONLY inside ``CookiePersistence.save``'s ``_save()``
closure, which runs on a worker thread via ``asyncio.to_thread``. It is never
held by an async context — a
blocking ``threading.Lock`` taken on the event-loop thread would stall every
other coroutine (keepalive, RPCs, cancellation) while a sibling worker thread
does file I/O. That is the priority-inversion failure mode this contract
exists to prevent.

These tests are deliberately structural: they verify the rule end-to-end
(via the live ``save_cookies`` path) AND statically (by scanning the cookie
persistence collaborator for any unexpected acquisition site), so a future
refactor that adds an event-loop acquisition fails fast in CI.
"""

from __future__ import annotations

import ast
import asyncio
import inspect
import threading
from pathlib import Path

import httpx
import pytest

from notebooklm._cookie_persistence import CookiePersistence
from notebooklm.auth import AuthTokens
from notebooklm.client import NotebookLMClient
from tests._helpers.client_factory import build_client_shell_for_tests


def _make_core(tmp_path: Path, *, cookie_saver=None) -> NotebookLMClient:
    """Build a minimal ``NotebookLMClient`` whose ``save_cookies`` is safe to call.

    Order matters: ``AuthTokens.__post_init__`` calls ``build_cookie_jar``,
    which loads from ``storage_path`` if it exists and enforces the cookie-set rule. We want it to take the in-memory ``cookies={...}``
    branch (file absent) so construction succeeds, THEN write the baseline
    file so the subsequent ``save_cookies`` call has something to merge
    against.

    ``cookie_saver`` (Phase 2 PR 4) is forwarded to ``NotebookLMClient(...)`` so
    tests can inject the persistence spy at construction rather than via
    the legacy ``notebooklm._core.save_cookies_to_storage`` monkeypatch.
    """
    storage_path = tmp_path / "storage_state.json"
    auth = AuthTokens(
        cookies={"SID": "x", "__Secure-1PSIDTS": "y"},
        csrf_token="t",
        session_id="s",
        storage_path=storage_path,
    )
    storage_path.write_text('{"cookies": []}')
    return build_client_shell_for_tests(auth, cookie_saver=cookie_saver)


@pytest.mark.asyncio
async def test_save_lock_acquired_off_event_loop_thread(
    tmp_path: Path,
) -> None:
    """The thread that holds ``CookiePersistence.save_lock`` MUST NOT be the loop thread.

    We spy on ``save_cookies_to_storage`` — which the production ``_save()``
    closure calls from inside ``with lock:`` — and record the thread it runs
    on. If a future refactor accidentally moves the ``with lock:``
    onto the loop thread (e.g. by inlining the closure into ``save_cookies``
    without ``asyncio.to_thread``), the spy will see the loop thread holding
    the lock, and this assertion will fail.
    """
    loop_thread = threading.current_thread()
    observed: dict[str, object] = {}

    # ``core`` is closed over by ``spy`` below; we declare a placeholder so
    # the spy's reference can be resolved before ``_make_core`` returns.
    core_ref: dict[str, NotebookLMClient] = {}

    def spy(jar, path, **kwargs):  # type: ignore[no-untyped-def]
        # ``save_cookies_to_storage`` is called from inside ``with lock:``
        # in ``_save()``. Whichever thread runs this spy is, by definition,
        # the thread currently holding ``CookiePersistence.save_lock``.
        observed["lock_held"] = core_ref[
            "core"
        ]._collaborators.cookie_persistence.save_lock.locked()
        observed["holder_thread"] = threading.current_thread()
        return True

    # Phase 2 PR 4: inject the cookie-saver seam via constructor injection
    # rather than via the legacy ``_core.save_cookies_to_storage`` string-target monkeypatch.
    core = _make_core(tmp_path, cookie_saver=spy)
    core_ref["core"] = core

    await core._collaborators.lifecycle.save_cookies(
        core._collaborators.cookie_persistence,
        httpx.Cookies(),
    )

    assert observed["lock_held"] is True, (
        "save_cookies must hold _save_lock for the duration of "
        "save_cookies_to_storage (precondition for the contract test)"
    )
    holder = observed["holder_thread"]
    assert isinstance(holder, threading.Thread)
    assert holder is not loop_thread, (
        "_save_lock contract violation: the lock was acquired on the "
        "event-loop thread. It must only be acquired inside the _save() "
        "closure dispatched via asyncio.to_thread, otherwise a blocking "
        "threading.Lock on the loop will stall every other coroutine "
        "(priority inversion)."
    )
    # Belt-and-braces: also compare ident in case some future Thread
    # subclass overrides ``__eq__``/``is`` semantics around object identity.
    # asyncio.to_thread dispatches onto the default ThreadPoolExecutor whose
    # workers are named ``asyncio_n``; we don't pin the exact name because
    # that's an implementation detail.
    assert holder.ident != loop_thread.ident


@pytest.mark.asyncio
async def test_save_lock_does_not_block_event_loop(
    tmp_path: Path,
) -> None:
    """While ``CookiePersistence.save_lock`` is held by a worker, the event loop must
    remain responsive.

    Direct proof of the no-priority-inversion property: hold the worker
    inside ``save_cookies_to_storage`` (which is called from inside
    ``with lock:``) and concurrently schedule loop work. If the
    loop were blocked on the lock, the heartbeat coroutine wouldn't run
    until the worker released; with the contract intact, the heartbeat
    observes the lock IS held while the loop is still scheduling.
    """
    in_save = threading.Event()
    release_save = threading.Event()
    loop_observations: list[bool] = []

    def spy(jar, path, **kwargs):  # type: ignore[no-untyped-def]
        in_save.set()
        # Hold the worker thread (and thus CookiePersistence.save_lock) until the loop has
        # demonstrated it can still schedule coroutines. Bounded to avoid a
        # hung test if the contract is ever violated and the loop deadlocks.
        assert release_save.wait(timeout=5.0), (
            "Loop never signalled release — likely event-loop blocked on "
            "_save_lock (contract violation)."
        )
        return True

    # Phase 2 PR 4: inject the cookie-saver seam at construction.
    core = _make_core(tmp_path, cookie_saver=spy)

    async def heartbeat() -> None:
        # Wait for the worker to enter the spy by polling — using asyncio.sleep
        # (not run_in_executor) so we don't contend for the default executor
        # that asyncio.to_thread also uses.
        for _ in range(500):  # ~5s ceiling at 10ms cadence
            if in_save.is_set():
                break
            await asyncio.sleep(0.01)
        # If we reach here, the loop is still scheduling tasks AND the
        # worker thread is inside the spy (lock held).
        loop_observations.append(core._collaborators.cookie_persistence.save_lock.locked())
        release_save.set()

    await asyncio.gather(
        core._collaborators.lifecycle.save_cookies(
            core._collaborators.cookie_persistence,
            httpx.Cookies(),
        ),
        heartbeat(),
    )

    assert loop_observations == [True], (
        "Event loop must remain responsive while _save_lock is held by a "
        "worker thread. If the loop blocks on the lock, heartbeat() can't "
        "observe `locked() is True` because it can't run at all — "
        f"observations={loop_observations!r}"
    )


def test_save_lock_only_acquired_inside_save_closure() -> None:
    """Static guard: the blocking lock is acquired inside the worker closure.

    This catches a refactor that adds a second ``with self.save_lock:`` or
    aliased ``with lock:`` elsewhere (e.g. inside an async method) before such
    a change can ship — static-only, so it has zero runtime cost and runs even
    when the async test infrastructure is offline.
    """

    source_path = Path(inspect.getsourcefile(CookiePersistence) or "")
    assert source_path.is_file(), f"could not locate _cookie_persistence.py source: {source_path!r}"
    tree = ast.parse(source_path.read_text())

    # Find every ``with self.save_lock:`` or closure-aliased ``with lock:`` by
    # walking the collaborator AST.
    acquisition_sites: list[tuple[str, int]] = []

    class _Visitor(ast.NodeVisitor):
        """Walk the module, tracking the enclosing function chain so any
        ``with lock:`` site can be attributed to the function
        that contains it (lets us check whether it sits inside ``_save``).
        """

        def __init__(self) -> None:
            self._enclosing_func: list[str] = []

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            self._enclosing_func.append(node.name)
            self.generic_visit(node)
            self._enclosing_func.pop()

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            self._enclosing_func.append(node.name)
            self.generic_visit(node)
            self._enclosing_func.pop()

        def _record_if_save_lock(self, expr: ast.expr) -> None:
            # Match ``self.save_lock`` directly.
            if (
                isinstance(expr, ast.Attribute)
                and expr.attr == "save_lock"
                and isinstance(expr.value, ast.Name)
                and expr.value.id == "self"
            ) or (isinstance(expr, ast.Name) and expr.id == "lock"):
                where = ".".join(self._enclosing_func) or "<module>"
                acquisition_sites.append((where, expr.lineno))

        def visit_With(self, node: ast.With) -> None:
            for item in node.items:
                self._record_if_save_lock(item.context_expr)
            self.generic_visit(node)

        def visit_AsyncWith(self, node: ast.AsyncWith) -> None:
            for item in node.items:
                # An ``async with CookiePersistence.save_lock:`` would itself be a contract
                # violation (lock is sync), so flag it the same way.
                self._record_if_save_lock(item.context_expr)
            self.generic_visit(node)

    _Visitor().visit(tree)

    offenders = [
        (where, lineno) for (where, lineno) in acquisition_sites if "_save" not in where.split(".")
    ]
    assert offenders == [], (
        "_save_lock contract violation: blocking lock acquisition found "
        f"outside the ``_save`` closure: {offenders!r}. The lock must "
        "ONLY be acquired inside ``_save()`` (run via asyncio.to_thread). "
        "See ``_cookie_persistence.py`` "
        "for the contract details."
    )
    assert acquisition_sites, "expected CookiePersistence.save to acquire the save lock"