"""Tests for ``notebooklm use`` fail-closed verification (PR). Before PR, ``notebooklm use `` persisted the supplied notebook ID to ``context.json`` *even when the existence check failed* — either because the RPC errored, or because the server returned a degenerate "empty notebook" payload for an unknown ID. The result was poisoned saved state that broke downstream commands until the user manually cleared the context. This module pins the post-fix contract: * Successful ``client.notebooks.get`` → context is persisted, exit 0. * ``NotebookNotFoundError`` from ``client.notebooks.get`` → exit 1, context file is NOT created. * ``--force`` flag → context is persisted regardless of verification. Each test mocks the entire client surface; nothing here hits the network. """ from __future__ import annotations import json from datetime import datetime from unittest.mock import AsyncMock, patch import pytest from click.testing import CliRunner import notebooklm.auth as auth_module import notebooklm.cli.context as context_module import notebooklm.cli.helpers as helpers_module import notebooklm.cli.resolve as resolve_module import notebooklm.cli.services.session_context as session_context_module import notebooklm.cli.session_cmd as session_cmd_module from notebooklm.exceptions import NotebookNotFoundError, RPCError from notebooklm.notebooklm_cli import cli from notebooklm.types import Notebook from .conftest import create_mock_client, inject_client @pytest.fixture def runner() -> CliRunner: return CliRunner() @pytest.fixture def mock_auth(): """Patch storage auth so the `use` command sees real-looking cookies.""" with patch.object(helpers_module, "load_auth_from_storage") as mock: mock.return_value = { "SID": "test", "__Secure-1PSIDTS": "test_1psidts", "HSID": "test", "SSID": "test", "APISID": "test", "SAPISID": "test", } yield mock @pytest.fixture def mock_context_file(tmp_path): """Provide a temporary context path; every consumer call site is patched. The ``notebooklm.cli.session_cmd.get_context_path`` re-export was retired in #1367; ``read_status`` now resolves the symbol on its real consumer module ``services.session_context`` (matching the canonical ``mock_context_file`` fixture in ``conftest.py``). """ context_file = tmp_path / "context.json" with ( patch.object(helpers_module, "get_context_path", return_value=context_file), patch.object(context_module, "get_context_path", return_value=context_file), patch.object(resolve_module, "get_context_path", return_value=context_file), patch.object( session_context_module, "get_context_path", return_value=context_file, ), ): yield context_file def _make_notebook(notebook_id: str = "nb_real", title: str = "Real Notebook") -> Notebook: return Notebook( id=notebook_id, title=title, created_at=datetime(2026, 1, 15), is_owner=True, ) class TestNotebookNotFoundIsRPCError: """``NotebookNotFoundError`` must still be catchable as ``RPCError``.""" def test_inherits_from_rpc_error(self): # ``except RPCError`` at higher layers must still match — this is the # whole point of widening the base class in the fail-closed fix. assert issubclass(NotebookNotFoundError, RPCError) def test_carries_notebook_id(self): err = NotebookNotFoundError("nb_typo") assert err.notebook_id == "nb_typo" assert "nb_typo" in str(err) def test_accepts_method_id(self): # The CLI/_notebooks.get site forwards method_id for diagnostics. err = NotebookNotFoundError("nb_x", method_id="rwIQyf") assert err.method_id == "rwIQyf" class TestUseFailsClosedBadId: """Bad notebook ID → no context write, exit 1.""" def test_notebook_not_found_does_not_persist(self, runner, mock_auth, mock_context_file): mock_client = create_mock_client() mock_client.notebooks.get = AsyncMock( side_effect=NotebookNotFoundError("nb_missing", method_id="rwIQyf"), ) with patch.object( auth_module, "fetch_tokens_with_domains", new_callable=AsyncMock ) as mock_fetch: mock_fetch.return_value = ("csrf", "session") with patch.object( session_cmd_module, "resolve_notebook_id", new_callable=AsyncMock ) as mock_resolve: mock_resolve.return_value = "nb_missing" result = runner.invoke(cli, ["use", "nb_missing"], obj=inject_client(mock_client)) # Exit non-zero, never touched context.json. assert result.exit_code == 1 assert not mock_context_file.exists() # Surface a clear "not found" message so the user can self-correct. assert "nb_missing" in result.output assert "not found" in result.output.lower() or "force" in result.output.lower() def test_generic_rpc_error_does_not_persist(self, runner, mock_auth, mock_context_file): """Network / RPC errors also fail closed — we can't confirm existence.""" mock_client = create_mock_client() mock_client.notebooks.get = AsyncMock( side_effect=RPCError("server hung up", method_id="rwIQyf"), ) with patch.object( auth_module, "fetch_tokens_with_domains", new_callable=AsyncMock ) as mock_fetch: mock_fetch.return_value = ("csrf", "session") with patch.object( session_cmd_module, "resolve_notebook_id", new_callable=AsyncMock ) as mock_resolve: mock_resolve.return_value = "nb_rpc_fail" result = runner.invoke(cli, ["use", "nb_rpc_fail"], obj=inject_client(mock_client)) assert result.exit_code == 1 assert not mock_context_file.exists() class TestUseFailsClosedGoodId: """Good notebook ID → context persisted, exit 0.""" def test_good_id_persists_context(self, runner, mock_auth, mock_context_file): mock_client = create_mock_client() mock_client.notebooks.get = AsyncMock( return_value=_make_notebook("nb_real", "Real Notebook"), ) with patch.object( auth_module, "fetch_tokens_with_domains", new_callable=AsyncMock ) as mock_fetch: mock_fetch.return_value = ("csrf", "session") with patch.object( session_cmd_module, "resolve_notebook_id", new_callable=AsyncMock ) as mock_resolve: mock_resolve.return_value = "nb_real" result = runner.invoke(cli, ["use", "nb_real"], obj=inject_client(mock_client)) assert result.exit_code == 0 assert mock_context_file.exists() data = json.loads(mock_context_file.read_text()) assert data["notebook_id"] == "nb_real" assert data["title"] == "Real Notebook" class TestUseForceFlag: """``--force`` bypasses the existence check entirely.""" def test_force_with_bad_id_still_persists(self, runner, mock_context_file): """No client mocking needed — --force never calls the network.""" result = runner.invoke(cli, ["use", "--force", "nb_offline"]) assert result.exit_code == 0 assert mock_context_file.exists() data = json.loads(mock_context_file.read_text()) assert data["notebook_id"] == "nb_offline" # Banner clarifies the ID was not verified, so the user isn't misled. assert "not verified" in result.output.lower() or "force" in result.output.lower() def test_force_does_not_call_get(self, runner, mock_auth, mock_context_file): """--force is offline-safe: even a guaranteed-raise ``get`` is bypassed.""" mock_client = create_mock_client() # If --force *did* call get, this mock would surface the error. mock_client.notebooks.get = AsyncMock( side_effect=NotebookNotFoundError("should-not-be-called"), ) result = runner.invoke( cli, ["use", "--force", "nb_force_id"], obj=inject_client(mock_client) ) assert result.exit_code == 0 assert mock_context_file.exists() # The get-mock side_effect would have surfaced if it had been invoked. mock_client.notebooks.get.assert_not_called()