teng-lin--notebooklm-py/tests/unit/cli/test_suggest_prompts.py

"""Tests for the ``notebooklm suggest-prompts`` CLI command.

Mirrors the sibling chat-command tests (``test_chat.py``): drives the command
through the same ``inject_client`` seam used by the top-level ``ask`` command,
patching ``client.notebooks.suggest_prompts`` rather than constructing a real
client. Covers the default mode, an explicit ``--mode``, ``--json`` envelope
shape, and an out-of-range ``--mode`` (the method's ``ValidationError`` must
surface as a clean exit-1 / VALIDATION_ERROR envelope, not a traceback).
"""

import json
from unittest.mock import AsyncMock, patch

import pytest
from click.testing import CliRunner

import notebooklm.auth as auth_module
from notebooklm.cli import helpers as helpers_module
from notebooklm.notebooklm_cli import cli
from notebooklm.types import PromptSuggestion

from .conftest import create_mock_client, inject_client

SUGGESTIONS = [
    PromptSuggestion(title="Professional Briefing", prompt="Give me a briefing on the sources."),
    PromptSuggestion(title="Key Risks", prompt="What are the key risks raised?"),
    PromptSuggestion(title="Counterarguments", prompt="What are the strongest counterarguments?"),
]


@pytest.fixture
def runner():
    return CliRunner()


@pytest.fixture
def mock_auth():
    with patch.object(helpers_module, "load_auth_from_storage") as mock:
        mock.return_value = {
            "SID": "test",
            "HSID": "test",
            "SSID": "test",
            "APISID": "test",
            "SAPISID": "test",
        }
        yield mock


def _invoke(runner, mock_client, argv):
    with patch.object(
        auth_module, "fetch_tokens_with_domains", new_callable=AsyncMock
    ) as mock_fetch:
        mock_fetch.return_value = ("csrf", "session")
        return runner.invoke(cli, argv, obj=inject_client(mock_client))


def test_default_mode_lists_suggestions(runner, mock_auth):
    mock_client = create_mock_client()
    mock_client.notebooks.suggest_prompts = AsyncMock(return_value=SUGGESTIONS)

    result = _invoke(runner, mock_client, ["suggest-prompts", "-n", "nb_123"])

    assert result.exit_code == 0, result.output
    # Default mode is 4 and source_ids defaults to None (all sources).
    mock_client.notebooks.suggest_prompts.assert_awaited_once()
    call = mock_client.notebooks.suggest_prompts.call_args
    assert call.kwargs["mode"] == 4
    assert call.kwargs["source_ids"] is None
    assert call.kwargs["query"] is None
    # Text output is a numbered list of title + prompt.
    assert "Professional Briefing" in result.output
    assert "Give me a briefing on the sources." in result.output
    assert "1." in result.output


def test_explicit_mode_and_query_forwarded(runner, mock_auth):
    mock_client = create_mock_client()
    mock_client.notebooks.suggest_prompts = AsyncMock(return_value=SUGGESTIONS)

    result = _invoke(
        runner,
        mock_client,
        ["suggest-prompts", "-n", "nb_123", "--mode", "8", "--query", "exam topics"],
    )

    assert result.exit_code == 0, result.output
    call = mock_client.notebooks.suggest_prompts.call_args
    assert call.kwargs["mode"] == 8
    assert call.kwargs["query"] == "exam topics"


def test_source_ids_forwarded(runner, mock_auth):
    mock_client = create_mock_client()
    mock_client.notebooks.suggest_prompts = AsyncMock(return_value=SUGGESTIONS)

    result = _invoke(
        runner,
        mock_client,
        ["suggest-prompts", "-n", "nb_123", "-s", "src_001", "-s", "src_002"],
    )

    assert result.exit_code == 0, result.output
    call = mock_client.notebooks.suggest_prompts.call_args
    assert call.kwargs["source_ids"] == ["src_001", "src_002"]


def test_json_envelope_shape(runner, mock_auth):
    mock_client = create_mock_client()
    mock_client.notebooks.suggest_prompts = AsyncMock(return_value=SUGGESTIONS)

    result = _invoke(runner, mock_client, ["suggest-prompts", "-n", "nb_123", "--json"])

    assert result.exit_code == 0, result.output
    payload = json.loads(result.stdout)
    assert payload["notebook_id"] == "nb_123"
    assert payload["count"] == 3
    assert payload["suggestions"][0] == {
        "title": "Professional Briefing",
        "prompt": "Give me a briefing on the sources.",
    }


def test_empty_suggestions_text(runner, mock_auth):
    mock_client = create_mock_client()
    mock_client.notebooks.suggest_prompts = AsyncMock(return_value=[])

    result = _invoke(runner, mock_client, ["suggest-prompts", "-n", "nb_123"])

    assert result.exit_code == 0, result.output
    assert "No prompt suggestions" in result.output


# ``11`` is the just-out-of-range boundary (cap is 10) — pins the CLI mirror so it
# can't silently widen past 10; ``99`` is a clearly-out value.
@pytest.mark.parametrize("bad_mode", ["11", "99"])
def test_bad_mode_exits_one(runner, mock_auth, bad_mode):
    mock_client = create_mock_client()
    mock_client.notebooks.suggest_prompts = AsyncMock(return_value=[])

    result = _invoke(
        runner,
        mock_client,
        ["suggest-prompts", "-n", "nb_123", "--mode", bad_mode, "--json"],
    )

    assert result.exit_code == 1, result.output
    payload = json.loads(result.stdout)
    assert payload["error"] is True
    assert payload["code"] == "VALIDATION_ERROR"
    assert "1..10" in payload["message"]


def test_bad_mode_validated_before_source_resolution(runner, mock_auth):
    # The mode check fires BEFORE notebook/source resolution, so an out-of-range
    # mode with an unresolvable partial ``-s`` id still surfaces the mode error
    # (not a "no source found" error) and never incurs a ``sources.list`` RPC.
    mock_client = create_mock_client()
    mock_client.sources.list = AsyncMock(return_value=[])
    mock_client.notebooks.suggest_prompts = AsyncMock(return_value=[])

    result = _invoke(
        runner,
        mock_client,
        ["suggest-prompts", "-n", "nb_123", "-s", "does-not-exist", "--mode", "99", "--json"],
    )

    assert result.exit_code == 1, result.output
    payload = json.loads(result.stdout)
    assert payload["code"] == "VALIDATION_ERROR"
    assert "1..10" in payload["message"]
    mock_client.sources.list.assert_not_awaited()
    mock_client.notebooks.suggest_prompts.assert_not_awaited()