Files
teng-lin--notebooklm-py/tests/unit/cli/test_research_characterization.py
wehub-resource-sync 09e9f3545f
Test / Code Quality (push) Has been cancelled
Test / Test (macos-latest, Python 3.10) (push) Has been cancelled
Test / Test (macos-latest, Python 3.11) (push) Has been cancelled
Test / Test (macos-latest, Python 3.12) (push) Has been cancelled
Test / Test (macos-latest, Python 3.13) (push) Has been cancelled
Test / Test (macos-latest, Python 3.14) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.10) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.11) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.12) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.13) (push) Has been cancelled
Test / Test (ubuntu-latest, Python 3.14) (push) Has been cancelled
Test / Test (windows-latest, Python 3.10) (push) Has been cancelled
Test / Test (windows-latest, Python 3.11) (push) Has been cancelled
Test / Test (windows-latest, Python 3.12) (push) Has been cancelled
Test / Test (windows-latest, Python 3.13) (push) Has been cancelled
Test / Test (windows-latest, Python 3.14) (push) Has been cancelled
CodeQL / Analyze (push) Has been cancelled
dependency-audit / pip-audit (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:30:13 +08:00

456 lines
17 KiB
Python

"""Golden-snapshot characterization tests for ``research wait``.
These tests freeze the current ``research wait`` behavior (text + JSON modes,
across the happy / timeout / cancelled / --no-import / --import-all paths)
before the P3.T6a extraction lands. They MUST pass on main HEAD (commit 1 of
the PR) and continue to pass byte-for-byte after the service extraction
(commit 2).
Covered paths (text + JSON):
* happy — research completes immediately, no import.
* timeout — research stays ``in_progress`` until ``--timeout`` elapses.
* cancelled — SIGINT during the polling spinner triggers the structured
cancellation envelope from ``emit_cancelled_and_exit``.
* --no-import — completed; ``--import-all`` NOT passed; importer is never
invoked (this is the default "wait without importing" shape
a script would use).
* --import-all — completed; ``--import-all`` calls ``import_with_retry``.
The JSON-mode snapshots compare parsed dicts (key set + values) so output is
stable across Click/Rich versions. The text-mode snapshots compare a small
normalized line slice (color codes stripped, fixed substrings only) so they
don't fight Rich rendering differences.
"""
from __future__ import annotations
import importlib
import json
import re
from unittest.mock import AsyncMock, patch
import pytest
from notebooklm.notebooklm_cli import cli
from .conftest import create_mock_client, inject_client, research_task
research_import_module = importlib.import_module("notebooklm.cli.research_import")
# Strip ANSI/Rich SGR sequences so text-mode snapshots survive Rich version
# bumps. We don't care that "[green]" or "\x1b[32m" is on/off — we care that
# the literal phrase appears.
_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
def _strip_ansi(text: str) -> str:
return _ANSI_RE.sub("", text)
@pytest.fixture
def runner_and_mocks(runner, mock_auth, mock_fetch_tokens):
"""Convenience fixture: returns ``(runner, mock_client_factory)``."""
def factory(poll_return):
mock_client = create_mock_client()
# ``research.poll`` now returns a typed ``ResearchTask``; adapt legacy
# dict specs declared by these characterization fixtures.
adapted = research_task(poll_return) if isinstance(poll_return, dict) else poll_return
mock_client.research.poll = AsyncMock(return_value=adapted)
return mock_client
return runner, factory
# ---------------------------------------------------------------------------
# Happy path — completed, no import
# ---------------------------------------------------------------------------
HAPPY_POLL = {
"status": "completed",
"task_id": "task_abc",
"query": "AI research",
"sources": [{"title": "Source 1", "url": "http://example.com"}],
"report": "# Report\nBody",
}
class TestWaitHappy:
def test_happy_text(self, runner_and_mocks):
runner, factory = runner_and_mocks
result = runner.invoke(
cli, ["research", "wait", "-n", "nb_123"], obj=inject_client(factory(HAPPY_POLL))
)
assert result.exit_code == 0
out = _strip_ansi(result.output)
# Golden lines (order-sensitive substrings).
assert "Research completed" in out
assert "AI research" in out
assert "Found 1 sources" in out
assert "Source 1" in out
assert "Report" in out
# Import was NOT invoked.
assert "Imported" not in out
def test_happy_json(self, runner_and_mocks):
runner, factory = runner_and_mocks
result = runner.invoke(
cli,
["research", "wait", "-n", "nb_123", "--json"],
obj=inject_client(factory(HAPPY_POLL)),
)
assert result.exit_code == 0
payload = json.loads(result.output)
# Frozen JSON contract (P1.T2 baseline).
assert payload == {
"status": "completed",
"query": "AI research",
"sources_found": 1,
# Typed sources serialize back to the canonical dict shape, which
# always carries ``result_type``.
"sources": [{"url": "http://example.com", "title": "Source 1", "result_type": 1}],
"report": "# Report\nBody",
}
# No import-related keys when --import-all was not passed.
assert "imported" not in payload
assert "cited_only" not in payload
# ---------------------------------------------------------------------------
# Timeout path — research never completes
# ---------------------------------------------------------------------------
TIMEOUT_POLL = {"status": "in_progress", "query": "AI research"}
class TestWaitTimeout:
def test_timeout_text(self, runner_and_mocks):
runner, factory = runner_and_mocks
result = runner.invoke(
cli,
["research", "wait", "-n", "nb_123", "--timeout", "1", "--interval", "1"],
obj=inject_client(factory(TIMEOUT_POLL)),
)
assert result.exit_code == 1
out = _strip_ansi(result.output)
assert "Timed out after 1 seconds" in out
def test_timeout_json(self, runner_and_mocks):
runner, factory = runner_and_mocks
result = runner.invoke(
cli,
[
"research",
"wait",
"-n",
"nb_123",
"--json",
"--timeout",
"1",
"--interval",
"1",
],
obj=inject_client(factory(TIMEOUT_POLL)),
)
assert result.exit_code == 1
payload = json.loads(result.output)
assert payload == {"status": "timeout", "error": "Timed out after 1s"}
# ---------------------------------------------------------------------------
# Cancelled path — SIGINT during polling triggers structured cancellation
# ---------------------------------------------------------------------------
class TestWaitCancelled:
"""Ctrl-C inside ``status_with_elapsed`` must surface the canonical
``Cancelled. Resume with: notebooklm research status`` shape.
Verifies the resume hint plumbed into ``status_with_elapsed`` matches the
one wired in ``research_cmd.research_wait`` today.
"""
def test_cancelled_text(self, runner_and_mocks):
runner, _ = runner_and_mocks
mock_client = create_mock_client()
mock_client.research.poll = AsyncMock(side_effect=KeyboardInterrupt)
result = runner.invoke(
cli, ["research", "wait", "-n", "nb_123"], obj=inject_client(mock_client)
)
# SIGINT exit code is 130 (128 + signal 2).
assert result.exit_code == 130
out = _strip_ansi(result.output)
# ``Cancelled.`` line + canonical resume hint.
assert "Cancelled" in out
assert "notebooklm research status" in out
def test_cancelled_json(self, runner_and_mocks):
runner, _ = runner_and_mocks
mock_client = create_mock_client()
mock_client.research.poll = AsyncMock(side_effect=KeyboardInterrupt)
result = runner.invoke(
cli, ["research", "wait", "-n", "nb_123", "--json"], obj=inject_client(mock_client)
)
assert result.exit_code == 130
payload = json.loads(result.output)
# Structured cancellation envelope from emit_cancelled_and_exit.
assert payload["error"] is True
assert payload["code"] == "CANCELLED"
# Resume hint surfaces the research-specific command.
assert "notebooklm research status" in payload.get("resume_hint", "")
# ---------------------------------------------------------------------------
# --no-import path — explicit "completed without importing" goldens
# ---------------------------------------------------------------------------
class TestWaitNoImport:
"""Default ``research wait`` (no ``--import-all``) returns the completed
payload but never calls the importer. This is the canonical script shape.
"""
def test_no_import_text_does_not_invoke_importer(self, runner_and_mocks):
runner, factory = runner_and_mocks
with patch.object(
research_import_module, "import_with_retry", new_callable=AsyncMock
) as mock_import:
result = runner.invoke(
cli, ["research", "wait", "-n", "nb_123"], obj=inject_client(factory(HAPPY_POLL))
)
assert result.exit_code == 0
mock_import.assert_not_awaited()
out = _strip_ansi(result.output)
assert "Research completed" in out
assert "Imported" not in out
def test_no_import_json_omits_import_keys(self, runner_and_mocks):
runner, factory = runner_and_mocks
with patch.object(
research_import_module, "import_with_retry", new_callable=AsyncMock
) as mock_import:
result = runner.invoke(
cli,
["research", "wait", "-n", "nb_123", "--json"],
obj=inject_client(factory(HAPPY_POLL)),
)
assert result.exit_code == 0
mock_import.assert_not_awaited()
payload = json.loads(result.output)
assert "imported" not in payload
assert "imported_sources" not in payload
assert "cited_only" not in payload
assert "cited_sources_selected" not in payload
assert "cited_only_fallback" not in payload
# ---------------------------------------------------------------------------
# --import-all path — completed + import + (optional) --cited-only
# ---------------------------------------------------------------------------
class TestWaitImportAll:
def test_import_all_text(self, runner_and_mocks):
runner, factory = runner_and_mocks
mock_client = factory(HAPPY_POLL)
with patch.object(
research_import_module, "import_with_retry", new_callable=AsyncMock
) as mock_import:
mock_import.return_value = [{"id": "src_1", "title": "Source 1"}]
result = runner.invoke(
cli,
["research", "wait", "-n", "nb_123", "--import-all"],
obj=inject_client(mock_client),
)
assert result.exit_code == 0
out = _strip_ansi(result.output)
assert "Research completed" in out
assert "Imported 1 sources" in out
# P1.T2 task-id pinning: the importer must receive the pinned task_id
# discovered on the first poll, not None.
mock_import.assert_awaited_once_with(
mock_client,
"nb_123",
"task_abc",
[{"url": "http://example.com", "title": "Source 1", "result_type": 1}],
max_elapsed=300,
)
def test_import_all_json(self, runner_and_mocks):
runner, factory = runner_and_mocks
mock_client = factory(HAPPY_POLL)
with patch.object(
research_import_module, "import_with_retry", new_callable=AsyncMock
) as mock_import:
mock_import.return_value = [{"id": "src_1", "title": "Source 1"}]
result = runner.invoke(
cli,
["research", "wait", "-n", "nb_123", "--json", "--import-all"],
obj=inject_client(mock_client),
)
assert result.exit_code == 0
payload = json.loads(result.output)
assert payload["status"] == "completed"
assert payload["sources_found"] == 1
assert payload["imported"] == 1
assert payload["imported_sources"] == [{"id": "src_1", "title": "Source 1"}]
# cited_only branch not active.
assert "cited_only" not in payload
mock_import.assert_awaited_once_with(
mock_client,
"nb_123",
"task_abc",
[{"url": "http://example.com", "title": "Source 1", "result_type": 1}],
max_elapsed=300,
json_output=True,
)
def test_import_all_cited_only_text(self, runner_and_mocks):
runner, _ = runner_and_mocks
mock_client = create_mock_client()
mock_client.research.poll = AsyncMock(
return_value=research_task(
{
"status": "completed",
"task_id": "task_abc",
"query": "AI research",
"sources": [
{"title": "Cited", "url": "https://example.com/cited"},
{"title": "Uncited", "url": "https://example.com/uncited"},
],
"report": "Report cites https://example.com/cited",
}
)
)
with patch.object(
research_import_module, "import_with_retry", new_callable=AsyncMock
) as mock_import:
mock_import.return_value = [{"id": "src_1", "title": "Cited"}]
result = runner.invoke(
cli,
["research", "wait", "-n", "nb_123", "--import-all", "--cited-only"],
obj=inject_client(mock_client),
)
assert result.exit_code == 0
out = _strip_ansi(result.output)
assert "Imported 1 sources" in out
# Only the cited source is forwarded to the importer.
mock_import.assert_awaited_once_with(
mock_client,
"nb_123",
"task_abc",
[{"url": "https://example.com/cited", "title": "Cited", "result_type": 1}],
max_elapsed=300,
)
def test_import_all_cited_only_json(self, runner_and_mocks):
runner, _ = runner_and_mocks
mock_client = create_mock_client()
mock_client.research.poll = AsyncMock(
return_value=research_task(
{
"status": "completed",
"task_id": "task_abc",
"query": "AI research",
"sources": [
{"title": "Cited", "url": "https://example.com/cited"},
{"title": "Uncited", "url": "https://example.com/uncited"},
],
"report": "Report cites https://example.com/cited",
}
)
)
with patch.object(
research_import_module, "import_with_retry", new_callable=AsyncMock
) as mock_import:
mock_import.return_value = [{"id": "src_1", "title": "Cited"}]
result = runner.invoke(
cli,
[
"research",
"wait",
"-n",
"nb_123",
"--json",
"--import-all",
"--cited-only",
],
obj=inject_client(mock_client),
)
assert result.exit_code == 0
payload = json.loads(result.output)
# P1.T2 fields preserved.
assert payload["cited_only"] is True
assert payload["cited_sources_selected"] == 1
assert payload["cited_only_fallback"] is False
assert payload["imported"] == 1
mock_import.assert_awaited_once_with(
mock_client,
"nb_123",
"task_abc",
[{"url": "https://example.com/cited", "title": "Cited", "result_type": 1}],
max_elapsed=300,
json_output=True,
)
# ---------------------------------------------------------------------------
# P1.T2 task-id pinning regression guard (must survive extraction unmodified)
# ---------------------------------------------------------------------------
class TestTaskIdPinning:
"""First poll discovers a task_id; subsequent polls MUST pin to it.
Repro of the P1.T2 fix: the wait loop pins the task_id discriminator after
the first non-empty discovery so a second concurrent research task cannot
substitute its sources/report into the in-flight wait.
"""
def test_task_id_pinned_after_first_discovery(self, runner_and_mocks):
runner, _ = runner_and_mocks
mock_client = create_mock_client()
# First poll: in_progress, with a task_id to pin. Second poll: completed.
poll_mock = AsyncMock()
poll_mock.side_effect = [
research_task({"status": "in_progress", "task_id": "task_pinned", "query": "AI"}),
research_task(
{
"status": "completed",
"task_id": "task_pinned",
"query": "AI",
"sources": [{"title": "S", "url": "http://example.com"}],
"report": "R",
}
),
]
mock_client.research.poll = poll_mock
result = runner.invoke(
cli,
["research", "wait", "-n", "nb_123", "--interval", "1", "--timeout", "10"],
obj=inject_client(mock_client),
)
assert result.exit_code == 0
# First call: task_id=None (nothing pinned yet).
first_call = poll_mock.await_args_list[0]
assert first_call.kwargs.get("task_id") is None
# Second call: task_id pinned to the value discovered on the first poll.
second_call = poll_mock.await_args_list[1]
assert second_call.kwargs.get("task_id") == "task_pinned"