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
1795 lines
75 KiB
Python
1795 lines
75 KiB
Python
"""Unit tests for SourcesAPI file upload pipeline and YouTube detection."""
|
|
|
|
import ast
|
|
import inspect
|
|
import textwrap
|
|
import warnings
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from urllib.parse import parse_qs, urlparse
|
|
|
|
import pytest
|
|
|
|
from notebooklm._source.upload import SourceUploadPipeline
|
|
from notebooklm._sources import SourcesAPI
|
|
from notebooklm.exceptions import NetworkError, RPCError, ValidationError
|
|
from notebooklm.rpc import RPCMethod
|
|
from notebooklm.rpc.types import SourceStatus
|
|
from notebooklm.types import Source
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_core():
|
|
"""Create a mocked Session for SourcesAPI."""
|
|
from tests._fixtures.fake_core import make_fake_core
|
|
|
|
core = make_fake_core(rpc_call=AsyncMock())
|
|
core.auth = MagicMock()
|
|
core.auth.authuser = 0
|
|
core.auth.account_email = None
|
|
# Upload paths pass the live http client's cookie jar to httpx so cookies
|
|
# are scoped by Domain attribute (#373). Keep this distinct from
|
|
# auth.cookie_jar so upload tests prove the live jar is used.
|
|
auth_cookie_jar = MagicMock(name="auth_cookie_jar")
|
|
live_cookie_jar = MagicMock(name="live_cookie_jar")
|
|
core.auth.cookie_jar = auth_cookie_jar
|
|
core.get_http_client = MagicMock()
|
|
core.get_http_client.return_value.cookies = live_cookie_jar
|
|
core.kernel = core
|
|
# The stateful upload pipeline reaches the live cookie jar through the
|
|
# injected kernel/session compatibility path. Keep this distinct from
|
|
# auth.cookie_jar so the invariant still proves uploads reuse the live jar.
|
|
core.live_cookies = MagicMock(return_value=live_cookie_jar)
|
|
# Mirror ``Session``'s auth-route helper surface. The
|
|
# ``authuser_query()`` and ``authuser_header()`` callables read
|
|
# ``core.auth.authuser`` / ``core.auth.account_email`` at call time so
|
|
# tests that mutate ``mock_core.auth.authuser`` mid-test still observe
|
|
# the updated routing.
|
|
from notebooklm._auth.account import authuser_query as _authuser_query
|
|
from notebooklm._auth.account import format_authuser_value as _authuser_header
|
|
|
|
core.authuser_query = MagicMock(
|
|
side_effect=lambda: _authuser_query(core.auth.authuser, core.auth.account_email)
|
|
)
|
|
core.authuser_header = MagicMock(
|
|
side_effect=lambda: _authuser_header(core.auth.authuser, core.auth.account_email)
|
|
)
|
|
core._drain_tracker = MagicMock()
|
|
core._drain_tracker.begin_transport_post = AsyncMock(return_value=object())
|
|
core._drain_tracker.finish_transport_post = AsyncMock()
|
|
core.operation_scope = MagicMock()
|
|
|
|
def operation_scope(_label):
|
|
@asynccontextmanager
|
|
async def scope() -> AsyncIterator[None]:
|
|
yield None
|
|
|
|
return scope()
|
|
|
|
core.operation_scope.side_effect = operation_scope
|
|
core.record_upload_queue_wait = MagicMock()
|
|
# Audit C1: :class:`SourceUploadPipeline` now takes ``LoopGuard``
|
|
# directly (the ``lifecycle`` constructor slot) so
|
|
# :meth:`SourceUploadPipeline.add_file` short-circuits cross-loop
|
|
# misuse. MagicMock blocks ``assert``-prefixed attribute access as a
|
|
# foot-gun guard, so the no-op stub must be installed explicitly.
|
|
core.assert_bound_loop = MagicMock()
|
|
return core
|
|
|
|
|
|
@pytest.fixture
|
|
def sources_api(mock_core):
|
|
"""Create SourcesAPI with mocked core.
|
|
|
|
The uploader is constructed explicitly from the same mocked core so the
|
|
upload-path tests still exercise the real :class:`SourceUploadPipeline`
|
|
while honoring the now-required ``uploader=`` kwarg on
|
|
:class:`SourcesAPI`. ``mock_core`` bundles ``rpc_call`` /
|
|
``operation_scope`` / ``assert_bound_loop`` so it structurally
|
|
satisfies all three of the pipeline's narrow collaborator slots.
|
|
"""
|
|
uploader = SourceUploadPipeline(
|
|
rpc=mock_core,
|
|
drain=mock_core,
|
|
lifecycle=mock_core,
|
|
kernel=mock_core.kernel,
|
|
auth=mock_core.auth,
|
|
record_upload_queue_wait=mock_core.record_upload_queue_wait,
|
|
)
|
|
return SourcesAPI(mock_core, uploader=uploader)
|
|
|
|
|
|
def test_sources_api_makes_uploader_share_lifecycle_collaborators(sources_api):
|
|
"""SourcesAPI is the single owner of the source-lifecycle verbs.
|
|
|
|
The upload pipeline's ``list_sources`` / ``get_source`` / ``wait_*``
|
|
helpers must delegate to the SAME ``SourceLister`` / ``SourcePoller``
|
|
instances the public API uses, rather than parallel copies built in the
|
|
pipeline constructor (issue #1205). ``SourcesAPI.__init__`` injects its
|
|
own collaborators via ``configure_source_lifecycle``.
|
|
"""
|
|
assert sources_api._uploader._lister is sources_api._lister
|
|
assert sources_api._uploader._poller is sources_api._poller
|
|
|
|
|
|
def _self_runtime_auth_attr_read(node: ast.AST, attr: str) -> bool:
|
|
return (
|
|
isinstance(node, ast.Attribute)
|
|
and node.attr == attr
|
|
and isinstance(node.value, ast.Attribute)
|
|
and node.value.attr == "auth"
|
|
and isinstance(node.value.value, ast.Attribute)
|
|
and node.value.value.attr == "_core"
|
|
and isinstance(node.value.value.value, ast.Name)
|
|
and node.value.value.value.id == "self"
|
|
)
|
|
|
|
|
|
def _self_core_http_client_cookies_read(node: ast.AST) -> bool:
|
|
return (
|
|
isinstance(node, ast.Attribute)
|
|
and node.attr == "cookies"
|
|
and isinstance(node.value, ast.Call)
|
|
and isinstance(node.value.func, ast.Attribute)
|
|
and node.value.func.attr == "get_http_client"
|
|
and isinstance(node.value.func.value, ast.Attribute)
|
|
and node.value.func.value.attr == "_core"
|
|
and isinstance(node.value.func.value.value, ast.Name)
|
|
and node.value.func.value.value.id == "self"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"helper",
|
|
[
|
|
SourcesAPI._start_resumable_upload,
|
|
SourcesAPI._upload_file_streaming,
|
|
SourcesAPI._cancel_upload_session,
|
|
SourceUploadPipeline.start_resumable_upload,
|
|
SourceUploadPipeline.upload_file_streaming,
|
|
SourceUploadPipeline.cancel_upload_session,
|
|
],
|
|
)
|
|
def test_upload_helpers_do_not_read_runtime_auth_or_live_cookies_directly(helper):
|
|
"""Upload helpers must route through the narrow capability Protocols, not broad core internals."""
|
|
tree = ast.parse(textwrap.dedent(inspect.getsource(helper)))
|
|
violations = [
|
|
ast.unparse(node)
|
|
for node in ast.walk(tree)
|
|
if _self_runtime_auth_attr_read(node, "authuser")
|
|
or _self_runtime_auth_attr_read(node, "account_email")
|
|
or _self_core_http_client_cookies_read(node)
|
|
]
|
|
|
|
assert violations == []
|
|
|
|
|
|
def _assert_async_client_uses_live_cookies(mock_client_cls, mock_core) -> None:
|
|
assert (
|
|
mock_client_cls.call_args.kwargs["cookies"]
|
|
is mock_core.get_http_client.return_value.cookies
|
|
)
|
|
assert mock_client_cls.call_args.kwargs["cookies"] is not mock_core.auth.cookie_jar
|
|
|
|
|
|
class TestWaitArgsKeywordOnly:
|
|
"""``wait`` / ``wait_timeout`` are keyword-only on the ``add_*`` methods.
|
|
|
|
The v0.4.x positional-wait compatibility shim was removed in v0.7.0
|
|
(see ``docs/deprecations.md``). Passing ``wait`` / ``wait_timeout``
|
|
positionally now raises ``TypeError`` from Python's own argument
|
|
binding rather than emitting a ``DeprecationWarning``.
|
|
"""
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("method_name", "args"),
|
|
[
|
|
("add_url", ("nb_123", "https://example.com/article")),
|
|
("add_text", ("nb_123", "Title", "Body")),
|
|
("add_file", ("nb_123", "document.pdf", None)),
|
|
(
|
|
"add_drive",
|
|
(
|
|
"nb_123",
|
|
"drive_file_id",
|
|
"Drive Doc",
|
|
"application/vnd.google-apps.document",
|
|
),
|
|
),
|
|
],
|
|
)
|
|
async def test_keyword_wait_args_forward(self, sources_api, method_name, args):
|
|
target = sources_api._uploader if method_name == "add_file" else sources_api._adder
|
|
patched = AsyncMock(return_value=Source(id="src_123", title="Source"))
|
|
setattr(target, method_name, patched)
|
|
|
|
method = getattr(sources_api, method_name)
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("error", DeprecationWarning)
|
|
result = await method(*args, wait=True, wait_timeout=45.0)
|
|
|
|
assert result.id == "src_123"
|
|
patched.assert_awaited_once()
|
|
assert patched.call_args.kwargs["wait"] is True
|
|
assert patched.call_args.kwargs["wait_timeout"] == 45.0
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("method_name", "args"),
|
|
[
|
|
("add_url", ("nb_123", "https://example.com/article")),
|
|
("add_text", ("nb_123", "Title", "Body")),
|
|
("add_file", ("nb_123", "document.pdf", None)),
|
|
(
|
|
"add_drive",
|
|
("nb_123", "drive_file_id", "Drive Doc", "application/vnd.google-apps.document"),
|
|
),
|
|
],
|
|
)
|
|
async def test_positional_wait_args_raise_type_error(self, sources_api, method_name, args):
|
|
method = getattr(sources_api, method_name)
|
|
with pytest.raises(TypeError):
|
|
await method(*args, True, 45.0)
|
|
|
|
|
|
# =============================================================================
|
|
# _extract_youtube_video_id() tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestExtractYoutubeVideoId:
|
|
"""Tests for YouTube video ID extraction."""
|
|
|
|
def test_extract_youtube_short_url(self, sources_api):
|
|
"""Test extraction from youtu.be short URLs."""
|
|
url = "https://youtu.be/dQw4w9WgXcQ"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result == "dQw4w9WgXcQ"
|
|
|
|
def test_extract_youtube_short_url_http(self, sources_api):
|
|
"""Test extraction from HTTP youtu.be URLs."""
|
|
url = "http://youtu.be/abc123_XYZ"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result == "abc123_XYZ"
|
|
|
|
def test_extract_youtube_standard_watch_url(self, sources_api):
|
|
"""Test extraction from standard watch URLs."""
|
|
url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result == "dQw4w9WgXcQ"
|
|
|
|
def test_extract_youtube_watch_url_no_www(self, sources_api):
|
|
"""Test extraction from watch URLs without www."""
|
|
url = "https://youtube.com/watch?v=abc123-_XY"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result == "abc123-_XY"
|
|
|
|
def test_extract_youtube_shorts_url(self, sources_api):
|
|
"""Test extraction from shorts URLs."""
|
|
url = "https://www.youtube.com/shorts/abc123DEF"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result == "abc123DEF"
|
|
|
|
def test_extract_youtube_shorts_url_no_www(self, sources_api):
|
|
"""Test extraction from shorts URLs without www."""
|
|
url = "https://youtube.com/shorts/xyz789"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result == "xyz789"
|
|
|
|
def test_extract_youtube_returns_none_for_non_youtube(self, sources_api):
|
|
"""Test that non-YouTube URLs return None."""
|
|
url = "https://example.com/video"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result is None
|
|
|
|
def test_extract_youtube_returns_none_for_invalid_format(self, sources_api):
|
|
"""Test that invalid YouTube URLs return None."""
|
|
url = "https://youtube.com/invalid/format"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result is None
|
|
|
|
def test_extract_youtube_with_hyphen_underscore_in_id(self, sources_api):
|
|
"""Test extraction with hyphens and underscores in video ID."""
|
|
url = "https://youtu.be/a-b_c-D_E-f"
|
|
result = sources_api._extract_youtube_video_id(url)
|
|
assert result == "a-b_c-D_E-f"
|
|
|
|
|
|
# =============================================================================
|
|
# _register_file_source() tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestRegisterFileSource:
|
|
"""Tests for file source registration."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_success(self, sources_api, mock_core):
|
|
"""Test successful file source registration.
|
|
|
|
The wrapper now captures a baseline source-list before the create,
|
|
so the happy path issues two RPCs: GET_NOTEBOOK (baseline) + the
|
|
ADD_SOURCE_FILE register. The baseline-diff is what prevents
|
|
mis-matching a pre-existing same-named source on a retry probe.
|
|
"""
|
|
# Response structure: [[[["source_id_123"]]]] - 4 levels with string at deepest
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["source_id_abc"]]]]
|
|
|
|
result = await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
assert result == "source_id_abc"
|
|
# 2 calls: baseline GET_NOTEBOOK + ADD_SOURCE_FILE register.
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 2
|
|
methods_called = [call.args[0] for call in mock_core.rpc_executor.rpc_call.await_args_list]
|
|
assert RPCMethod.ADD_SOURCE_FILE in methods_called
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_parses_deeply_nested(self, sources_api, mock_core):
|
|
"""Test parsing deeply nested response."""
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["my_source_id"]]]]
|
|
|
|
result = await sources_api._register_file_source("nb_123", "doc.docx")
|
|
|
|
assert result == "my_source_id"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_raises_on_null_response(self, sources_api, mock_core):
|
|
"""Test that null response raises SourceAddError."""
|
|
from notebooklm.exceptions import SourceAddError
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = None
|
|
|
|
with pytest.raises(SourceAddError, match="Failed to get SOURCE_ID"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_raises_on_empty_response(self, sources_api, mock_core):
|
|
"""Test that empty response raises SourceAddError."""
|
|
from notebooklm.exceptions import SourceAddError
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = []
|
|
|
|
with pytest.raises(SourceAddError, match="Failed to get SOURCE_ID"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_extracts_id_from_nested_lists(self, sources_api, mock_core):
|
|
"""Test that ID is extracted from legacy singleton list envelopes."""
|
|
mock_core.rpc_executor.rpc_call.return_value = [[["source_id_123"]]]
|
|
|
|
result = await sources_api._register_file_source("nb_123", "test.pdf")
|
|
assert result == "source_id_123"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_raises_on_non_string_id(self, sources_api, mock_core):
|
|
"""Test that non-string source ID raises SourceAddError."""
|
|
from notebooklm.exceptions import SourceAddError
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[[[[12345]]]]]]
|
|
|
|
with pytest.raises(SourceAddError, match="Failed to get SOURCE_ID"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_handles_leading_none_shape(self, sources_api, mock_core):
|
|
"""The trusted prefixed-singleton envelope path should unwrap the
|
|
UUID-shaped SOURCE_ID after the leading None.
|
|
"""
|
|
uuid = "dc84ca28-2629-49ac-aec3-de45f0ec93e4"
|
|
mock_core.rpc_executor.rpc_call.return_value = [None, [[[uuid]]]]
|
|
|
|
result = await sources_api._register_file_source("nb_123", "report.pdf")
|
|
assert result == uuid
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"filename,response,expected",
|
|
[
|
|
# Filename echoed next to the SOURCE_ID (sibling).
|
|
(
|
|
"report.pdf",
|
|
[[[["report.pdf", ["11111111-2222-3333-4444-555555555555"]]]]],
|
|
"11111111-2222-3333-4444-555555555555",
|
|
),
|
|
# Filename at the first leaf position — the legacy data[0] walker
|
|
# would have returned 'doc.pdf' and broken downstream upload (#474).
|
|
(
|
|
"doc.pdf",
|
|
[[["doc.pdf"], [["aabbccdd-eeff-1122-3344-556677889900"]]]],
|
|
"aabbccdd-eeff-1122-3344-556677889900",
|
|
),
|
|
],
|
|
ids=["filename-sibling-of-uuid", "filename-at-position-zero"],
|
|
)
|
|
async def test_register_file_source_prefers_uuid_over_echoed_filename(
|
|
self, sources_api, mock_core, filename, response, expected
|
|
):
|
|
"""The extractor must skip the echoed filename and return the paired UUID.
|
|
|
|
The filename provides context for the SOURCE_ID (#474).
|
|
"""
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
[["", []]],
|
|
response,
|
|
]
|
|
|
|
result = await sources_api._register_file_source("nb_123", filename)
|
|
assert result == expected
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_falls_back_to_non_uuid_string(self, sources_api, mock_core):
|
|
"""Legacy singleton envelopes still accept plausible id-like non-UUID strings.
|
|
|
|
This applies when no UUID candidate is present.
|
|
"""
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_pdf"]]]]
|
|
|
|
result = await sources_api._register_file_source("nb_123", "doc.pdf")
|
|
assert result == "src_pdf"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_error_message_includes_shape_preview(
|
|
self, sources_api, mock_core
|
|
):
|
|
"""Future shape drift should report a sanitized response shape."""
|
|
from notebooklm.exceptions import SourceAddError
|
|
|
|
# Pure-numeric response — no string leaves → no candidates → raises.
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[1, 2, 3]]]
|
|
|
|
with pytest.raises(
|
|
SourceAddError,
|
|
match="Failed to get SOURCE_ID: no trustworthy SOURCE_ID found in array",
|
|
) as exc_info:
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
assert "[[[1, 2, 3]]]" not in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("status_token", ["OK", "DONE", "true", "null"])
|
|
async def test_register_file_source_rejects_all_alpha_status_tokens(
|
|
self, sources_api, mock_core, status_token
|
|
):
|
|
"""Fallback must reject all-alpha status tokens — ``OK``/``DONE``/``true``
|
|
in the response position would otherwise be picked up as a bogus
|
|
SOURCE_ID and break the downstream upload silently.
|
|
"""
|
|
from notebooklm.exceptions import SourceAddError
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[status_token]]]
|
|
|
|
with pytest.raises(SourceAddError, match="Failed to get SOURCE_ID"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_wraps_rpc_error_with_status_code(
|
|
self, sources_api, mock_core
|
|
):
|
|
"""When the decoder raises RPCError (e.g. server returned null result
|
|
with a status code at wrb.fr[5] — the suspected #474 mode for account-
|
|
routing mismatches), wrap the rich diagnostic into SourceAddError so
|
|
the user sees actionable detail instead of a generic "Failed to get
|
|
SOURCE_ID" with no context.
|
|
"""
|
|
from notebooklm.exceptions import ClientError, SourceAddError
|
|
|
|
# ClientError is what the decoder raises for status codes 5/7 (NOT_FOUND
|
|
# / PERMISSION_DENIED) with an account-routing hint attached — exactly
|
|
# the #114/#294 pattern we suspect for #474.
|
|
mock_core.rpc_executor.rpc_call.side_effect = ClientError(
|
|
"RPC o4cbdc returned null result with status code 7 (Permission denied). "
|
|
"If you have multiple Google accounts signed in...",
|
|
method_id="o4cbdc",
|
|
rpc_code=7,
|
|
)
|
|
|
|
with pytest.raises(SourceAddError, match="Permission denied") as exc_info:
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
# The original RPCError is preserved as the cause so debuggers can
|
|
# inspect the full decoder context.
|
|
assert isinstance(exc_info.value.cause, ClientError)
|
|
assert exc_info.value.cause.rpc_code == 7
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_lets_auth_error_propagate(self, sources_api, mock_core):
|
|
"""AuthError must NOT be wrapped — the core auth-refresh retry path
|
|
relies on the exception bubbling up unchanged.
|
|
"""
|
|
from notebooklm.exceptions import AuthError
|
|
|
|
mock_core.rpc_executor.rpc_call.side_effect = AuthError("session expired")
|
|
|
|
with pytest.raises(AuthError, match="session expired"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_lets_rate_limit_error_propagate(
|
|
self, sources_api, mock_core
|
|
):
|
|
"""RateLimitError must propagate so callers implementing back-off keep
|
|
their specific exception type (and ``retry_after``) without having to
|
|
unwrap ``SourceAddError.cause``.
|
|
"""
|
|
from notebooklm.exceptions import RateLimitError
|
|
|
|
mock_core.rpc_executor.rpc_call.side_effect = RateLimitError(
|
|
"API rate limit exceeded",
|
|
method_id="o4cbdc",
|
|
rpc_code="USER_DISPLAYABLE_ERROR",
|
|
)
|
|
|
|
with pytest.raises(RateLimitError, match="rate limit"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_lets_server_error_propagate(self, sources_api, mock_core):
|
|
"""ServerError must propagate so callers handling transient 5xx
|
|
backend errors keep the specific exception type for retry logic.
|
|
"""
|
|
from notebooklm.exceptions import ServerError
|
|
|
|
mock_core.rpc_executor.rpc_call.side_effect = ServerError(
|
|
"Backend unavailable",
|
|
method_id="o4cbdc",
|
|
rpc_code=500,
|
|
)
|
|
|
|
with pytest.raises(ServerError, match="Backend unavailable"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_file_source_walker_has_recursion_guard(self, sources_api, mock_core):
|
|
"""A deeply-nested response must not trigger RecursionError.
|
|
|
|
The depth guard caps traversal at ``_SOURCE_ID_ENVELOPE_MAX_DEPTH``.
|
|
"""
|
|
# 200-deep nest with a UUID at the bottom. Past the depth guard the
|
|
# walker stops, so the UUID is unreachable and we raise — but we don't
|
|
# crash with RecursionError.
|
|
from notebooklm.exceptions import SourceAddError
|
|
|
|
deep: list = ["dc84ca28-2629-49ac-aec3-de45f0ec93e4"]
|
|
for _ in range(200):
|
|
deep = [deep]
|
|
mock_core.rpc_executor.rpc_call.return_value = deep
|
|
|
|
with pytest.raises(SourceAddError, match="Failed to get SOURCE_ID"):
|
|
await sources_api._register_file_source("nb_123", "test.pdf")
|
|
|
|
|
|
# =============================================================================
|
|
# _start_resumable_upload() tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestStartResumableUpload:
|
|
"""Tests for starting resumable upload."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_resumable_upload_success(self, sources_api, mock_core):
|
|
"""Test successful upload start."""
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session123"
|
|
}
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api._start_resumable_upload(
|
|
"nb_123", "test.pdf", 1024, "src_456", "application/pdf"
|
|
)
|
|
|
|
assert result == "https://notebooklm.google.com/upload/_/?upload_id=session123"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_resumable_upload_includes_correct_headers(self, sources_api, mock_core):
|
|
"""Test that upload start includes correct headers."""
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._start_resumable_upload(
|
|
"nb_123", "test.pdf", 2048, "src_789", "application/pdf"
|
|
)
|
|
|
|
call_kwargs = mock_client.post.call_args[1]
|
|
headers = call_kwargs["headers"]
|
|
|
|
assert headers["x-goog-upload-command"] == "start"
|
|
assert headers["x-goog-upload-header-content-length"] == "2048"
|
|
assert headers["x-goog-upload-header-content-type"] == "application/pdf"
|
|
assert headers["x-goog-upload-protocol"] == "resumable"
|
|
# Cookie header is no longer set manually; httpx scopes cookies
|
|
# by Domain attribute via the cookie_jar kwarg (#373).
|
|
assert "Cookie" not in headers
|
|
_assert_async_client_uses_live_cookies(mock_client_cls, mock_core)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_resumable_upload_uses_integer_authuser_query_and_header(
|
|
self, sources_api, mock_core
|
|
):
|
|
"""Integer selected-account routing is preserved in URL and header."""
|
|
mock_core.auth.authuser = 2
|
|
mock_core.auth.account_email = None
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._start_resumable_upload(
|
|
"nb_123", "test.pdf", 2048, "src_789", "application/pdf"
|
|
)
|
|
|
|
url = mock_client.post.call_args.args[0]
|
|
assert parse_qs(urlparse(url).query) == {"authuser": ["2"]}
|
|
assert mock_client.post.call_args.kwargs["headers"]["x-goog-authuser"] == "2"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_resumable_upload_uses_encoded_account_email_query_and_header(
|
|
self, sources_api, mock_core
|
|
):
|
|
"""Account email wins over integer routing and is URL-encoded in the query."""
|
|
mock_core.auth.authuser = 2
|
|
mock_core.auth.account_email = "user+test@example.com"
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._start_resumable_upload(
|
|
"nb_123", "test.pdf", 2048, "src_789", "application/pdf"
|
|
)
|
|
|
|
url = mock_client.post.call_args.args[0]
|
|
assert "authuser=user%2Btest%40example.com" in urlparse(url).query
|
|
assert parse_qs(urlparse(url).query) == {"authuser": ["user+test@example.com"]}
|
|
assert (
|
|
mock_client.post.call_args.kwargs["headers"]["x-goog-authuser"]
|
|
== "user+test@example.com"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_resumable_upload_includes_json_body(self, sources_api, mock_core):
|
|
"""Test that upload start includes correct JSON body."""
|
|
import json
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._start_resumable_upload(
|
|
"nb_test", "myfile.pdf", 1000, "src_abc", "application/pdf"
|
|
)
|
|
|
|
call_kwargs = mock_client.post.call_args[1]
|
|
body = json.loads(call_kwargs["content"])
|
|
|
|
assert body["PROJECT_ID"] == "nb_test"
|
|
assert body["SOURCE_NAME"] == "myfile.pdf"
|
|
assert body["SOURCE_ID"] == "src_abc"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_resumable_upload_raises_on_missing_url_header(
|
|
self, sources_api, mock_core
|
|
):
|
|
"""Test that missing upload URL header raises SourceAddError."""
|
|
from notebooklm.exceptions import SourceAddError
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.headers = {} # No x-goog-upload-url
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
with pytest.raises(SourceAddError, match="Failed to get upload URL"):
|
|
await sources_api._start_resumable_upload(
|
|
"nb_123", "test.pdf", 1024, "src_456", "application/pdf"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_resumable_upload_raises_on_http_error(self, sources_api, mock_core):
|
|
"""Test that HTTP error raises exception."""
|
|
import httpx
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = httpx.HTTPStatusError(
|
|
"Server Error", request=MagicMock(), response=MagicMock()
|
|
)
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
await sources_api._start_resumable_upload(
|
|
"nb_123", "test.pdf", 1024, "src_456", "application/pdf"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# _upload_file_streaming() tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestUploadFileStreaming:
|
|
"""Tests for streaming file upload."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_file_streaming_success(self, sources_api, mock_core, tmp_path):
|
|
"""Test successful streaming file upload."""
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_bytes(b"file content here")
|
|
mock_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
# Should not raise
|
|
await sources_api._upload_file_streaming(
|
|
"https://notebooklm.google.com/upload/_/?upload_id=session", test_file
|
|
)
|
|
|
|
mock_client.post.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_file_streaming_includes_correct_headers(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""Test that streaming upload includes correct headers."""
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_bytes(b"content")
|
|
mock_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._upload_file_streaming(
|
|
"https://notebooklm.google.com/upload/_/?upload_id=session", test_file
|
|
)
|
|
|
|
call_kwargs = mock_client.post.call_args[1]
|
|
headers = call_kwargs["headers"]
|
|
|
|
assert headers["x-goog-upload-command"] == "upload, finalize"
|
|
assert headers["x-goog-upload-offset"] == "0"
|
|
# Cookie header is no longer set manually; httpx scopes cookies
|
|
# by Domain attribute via the cookie_jar kwarg (#373).
|
|
assert "Cookie" not in headers
|
|
_assert_async_client_uses_live_cookies(mock_client_cls, mock_core)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_file_streaming_preserves_authuser_header(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""Finalize upload sends the same selected-account header value."""
|
|
mock_core.auth.authuser = 2
|
|
mock_core.auth.account_email = "user+test@example.com"
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_bytes(b"content")
|
|
mock_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._upload_file_streaming(
|
|
"https://notebooklm.google.com/upload/_/?upload_id=session", test_file
|
|
)
|
|
|
|
assert (
|
|
mock_client.post.call_args.kwargs["headers"]["x-goog-authuser"]
|
|
== "user+test@example.com"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_upload_session_preserves_authuser_header_and_live_cookies(
|
|
self, sources_api, mock_core
|
|
):
|
|
"""Best-effort Scotty cancel keeps the provided auth route and live cookie jar."""
|
|
auth_route = "user+test@example.com"
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._cancel_upload_session(
|
|
"https://notebooklm.google.com/upload/_/?upload_id=session",
|
|
"https://notebooklm.google.com",
|
|
auth_route,
|
|
)
|
|
|
|
headers = mock_client.post.call_args.kwargs["headers"]
|
|
assert headers["x-goog-authuser"] == auth_route
|
|
assert headers["x-goog-upload-command"] == "cancel"
|
|
_assert_async_client_uses_live_cookies(mock_client_cls, mock_core)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_file_streaming_uses_generator(self, sources_api, mock_core, tmp_path):
|
|
"""Test that file content is streamed via generator."""
|
|
test_file = tmp_path / "test.txt"
|
|
test_content = b"This is my file content"
|
|
test_file.write_bytes(test_content)
|
|
mock_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.return_value = mock_response
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api._upload_file_streaming(
|
|
"https://notebooklm.google.com/upload/_/?upload_id=session", test_file
|
|
)
|
|
|
|
call_kwargs = mock_client.post.call_args[1]
|
|
# Content should be a generator, not bytes
|
|
content = call_kwargs["content"]
|
|
# Consume the generator to verify it yields the file content
|
|
chunks = [chunk async for chunk in content]
|
|
assert b"".join(chunks) == test_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_file_streaming_raises_on_http_error(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""Test that HTTP error raises exception."""
|
|
import httpx
|
|
|
|
test_file = tmp_path / "test.txt"
|
|
test_file.write_bytes(b"content")
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = httpx.HTTPStatusError(
|
|
"Upload Failed", request=MagicMock(), response=MagicMock()
|
|
)
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
await sources_api._upload_file_streaming(
|
|
"https://notebooklm.google.com/upload/_/?upload_id=session", test_file
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# add_file() tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestAddFile:
|
|
"""Tests for the add_file() public method."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_complete_flow(self, sources_api, mock_core, tmp_path):
|
|
"""Test complete file upload flow."""
|
|
# Create a temp file
|
|
test_file = tmp_path / "test.pdf"
|
|
test_file.write_bytes(b"fake pdf content")
|
|
|
|
# Mock the registration response - 4 levels with string at deepest
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_new_123"]]]]
|
|
|
|
# Mock HTTP calls
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file("nb_123", str(test_file))
|
|
|
|
assert result.id == "src_new_123"
|
|
assert result.title == "test.pdf"
|
|
assert result.kind == "unknown"
|
|
# 2 RPCs: GET_NOTEBOOK baseline + ADD_SOURCE_FILE register.
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 2
|
|
mock_core.operation_scope.assert_called_once_with("upload:0")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_raises_file_not_found(self, sources_api, mock_core):
|
|
"""Test that non-existent file raises FileNotFoundError."""
|
|
with pytest.raises(FileNotFoundError, match="File not found"):
|
|
await sources_api.add_file("nb_123", "/nonexistent/path/file.pdf")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_with_path_object(self, sources_api, mock_core, tmp_path):
|
|
"""Test add_file accepts Path objects."""
|
|
test_file = tmp_path / "doc.txt"
|
|
test_file.write_bytes(b"text content")
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_txt"]]]]
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file("nb_123", test_file) # Path object
|
|
|
|
assert result.id == "src_txt"
|
|
assert result.title == "doc.txt"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_keyword_wait_args(self, sources_api, mock_core, tmp_path):
|
|
"""``wait`` / ``wait_timeout`` keyword callers forward through to the wait."""
|
|
test_file = tmp_path / "report.pdf"
|
|
test_file.write_bytes(b"fake pdf content")
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_pdf"]]]]
|
|
sources_api._uploader.wait_until_ready = AsyncMock(
|
|
return_value=MagicMock(id="src_pdf", title="report.pdf")
|
|
)
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file(
|
|
"nb_123", str(test_file), None, wait=True, wait_timeout=45.0
|
|
)
|
|
|
|
assert result.id == "src_pdf"
|
|
sources_api._uploader.wait_until_ready.assert_awaited_once_with(
|
|
"nb_123", "src_pdf", timeout=45.0, transient_error_types=()
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_mime_type_non_none_sets_upload_content_type(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""Passing ``mime_type`` sets the Scotty upload content-type header."""
|
|
test_file = tmp_path / "image.png"
|
|
test_file.write_bytes(b"\x89PNG\r\n\x1a\n")
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_png"]]]]
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
result = await sources_api.add_file("nb_123", str(test_file), mime_type="image/png")
|
|
|
|
assert result.id == "src_png"
|
|
assert not [w for w in caught if issubclass(w.category, DeprecationWarning)]
|
|
start_headers = mock_client.post.await_args_list[0].kwargs["headers"]
|
|
assert start_headers["x-goog-upload-header-content-type"] == "image/png"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_mime_type_none_does_not_warn(self, sources_api, mock_core, tmp_path):
|
|
"""``mime_type=None`` (the default) MUST NOT emit ``DeprecationWarning``.
|
|
|
|
Default-path callers — the overwhelming majority — should never see the
|
|
deprecation message; otherwise it becomes noise rather than a signal.
|
|
"""
|
|
test_file = tmp_path / "notes.txt"
|
|
test_file.write_bytes(b"hello")
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_default"]]]]
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
result = await sources_api.add_file("nb_123", str(test_file))
|
|
|
|
deprecation_warnings = [w for w in caught if issubclass(w.category, DeprecationWarning)]
|
|
assert not deprecation_warnings, (
|
|
f"Unexpected DeprecationWarning emitted for mime_type=None default path: "
|
|
f"{[str(w.message) for w in deprecation_warnings]}"
|
|
)
|
|
assert result.id == "src_default"
|
|
start_headers = mock_client.post.await_args_list[0].kwargs["headers"]
|
|
assert start_headers["x-goog-upload-header-content-type"] == "text/plain"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_with_custom_title_renames_after_upload(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""A non-null title that differs from the filename should trigger a rename
|
|
and surface in the returned Source (regression test for #313).
|
|
|
|
Per #388, supplying a title forces a brief wait before the rename so the
|
|
UPDATE_SOURCE RPC reaches a registered source instead of silently
|
|
no-opping.
|
|
"""
|
|
test_file = tmp_path / "boring-filename.md"
|
|
test_file.write_bytes(b"# content\n")
|
|
|
|
# 3 rpc_call invocations in order:
|
|
# [0] baseline GET_NOTEBOOK for the probe-then-create wrapper
|
|
# [1] ADD_SOURCE_FILE register
|
|
# [2] UPDATE_SOURCE rename
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None, # baseline returns no useful list — empty notebook
|
|
[[[["src_md"]]]],
|
|
[[[["src_md"], "Real Intended Title", [None, None, None, None, 8]]]],
|
|
]
|
|
# The forced pre-rename registration wait is mocked — we don't want
|
|
# this test to depend on the polling implementation. It returns the
|
|
# registered source so add_file then issues the rename RPC.
|
|
sources_api._uploader.wait_until_registered = AsyncMock(
|
|
return_value=Source(id="src_md", title="boring-filename.md", _type_code=8)
|
|
)
|
|
sources_api._uploader.wait_until_ready = AsyncMock()
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file(
|
|
"nb_123", str(test_file), title=" Real Intended Title "
|
|
)
|
|
|
|
assert result.id == "src_md"
|
|
assert result.title == "Real Intended Title"
|
|
# 1 baseline GET_NOTEBOOK + 1 register + 1 rename
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 3
|
|
rename_params = mock_core.rpc_executor.rpc_call.call_args_list[2].args[1]
|
|
assert rename_params == [None, ["src_md"], [[["Real Intended Title"]]]]
|
|
# Narrow wait uses the caller's wait_timeout (default 120s) — not the
|
|
# full wait_until_ready. wait_until_registered returns on first
|
|
# PROCESSING/READY status so the bound stays cheap in practice.
|
|
sources_api._uploader.wait_until_registered.assert_awaited_once_with(
|
|
"nb_123", "src_md", timeout=120.0, transient_error_types=()
|
|
)
|
|
sources_api._uploader.wait_until_ready.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_with_custom_title_renames_after_wait(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""The custom-title rename must run AFTER the source is fully registered
|
|
server-side, otherwise the UPDATE_SOURCE RPC silently no-ops (#388).
|
|
With wait=True the order is: register -> upload -> wait -> rename.
|
|
"""
|
|
test_file = tmp_path / "boring-filename.md"
|
|
test_file.write_bytes(b"# content\n")
|
|
|
|
# 3 rpc_call invocations: baseline + register + rename (see the
|
|
# earlier test for the same pattern).
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None,
|
|
[[[["src_md"]]]],
|
|
[[[["src_md"], "Real Intended Title"]]],
|
|
]
|
|
|
|
async def wait_side_effect(notebook_id, source_id, *, timeout, transient_error_types):
|
|
assert notebook_id == "nb_123"
|
|
assert source_id == "src_md"
|
|
assert timeout == 120.0
|
|
assert transient_error_types == ()
|
|
# Wait must happen BEFORE the rename: at this point only the
|
|
# baseline GET_NOTEBOOK + ADD_SOURCE_FILE register RPCs have
|
|
# fired (2 total — see register_file_source's probe-then-create
|
|
# design for why the baseline call is required).
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 2
|
|
return Source(id=source_id, title="boring-filename.md", _type_code=8)
|
|
|
|
sources_api._uploader.wait_until_ready = AsyncMock(side_effect=wait_side_effect)
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file(
|
|
"nb_123",
|
|
str(test_file),
|
|
title="Real Intended Title",
|
|
wait=True,
|
|
)
|
|
|
|
assert result.title == "Real Intended Title"
|
|
sources_api._uploader.wait_until_ready.assert_awaited_once_with(
|
|
"nb_123", "src_md", timeout=120.0, transient_error_types=()
|
|
)
|
|
# 3 RPCs in total: baseline + register + rename.
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_with_title_forces_wait_when_wait_false(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""Even with wait=False, supplying a custom title must force a brief wait
|
|
so the rename hits a registered source instead of silently no-opping (#388).
|
|
"""
|
|
test_file = tmp_path / "boring-filename.md"
|
|
test_file.write_bytes(b"# content\n")
|
|
|
|
# 3 RPCs: baseline + register + rename.
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None,
|
|
[[[["src_md"]]]],
|
|
[[[["src_md"], "Real Intended Title", [None, None, None, None, 8]]]],
|
|
]
|
|
|
|
async def wait_side_effect(notebook_id, source_id, *, timeout, transient_error_types):
|
|
# Narrow registration wait — wait_timeout default (120s) is forwarded
|
|
# verbatim. wait_until_registered returns on the first non-PREPARING
|
|
# status, so the bound stays cheap.
|
|
assert timeout == 120.0
|
|
assert transient_error_types == ()
|
|
# Wait runs BEFORE the rename: at this point only the baseline
|
|
# GET_NOTEBOOK + ADD_SOURCE_FILE register RPCs have fired (2
|
|
# total).
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 2
|
|
return Source(id=source_id, title="boring-filename.md", _type_code=8)
|
|
|
|
sources_api._uploader.wait_until_registered = AsyncMock(side_effect=wait_side_effect)
|
|
sources_api._uploader.wait_until_ready = AsyncMock()
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file(
|
|
"nb_123",
|
|
str(test_file),
|
|
title="Real Intended Title",
|
|
)
|
|
|
|
assert result.title == "Real Intended Title"
|
|
sources_api._uploader.wait_until_registered.assert_awaited_once_with(
|
|
"nb_123", "src_md", timeout=120.0, transient_error_types=()
|
|
)
|
|
sources_api._uploader.wait_until_ready.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_with_title_forwards_wait_timeout(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""The caller's wait_timeout must be forwarded verbatim to the narrow
|
|
registration wait. The wait_until_registered helper polls and returns
|
|
on first PROCESSING/READY response, so generous bounds (e.g. long
|
|
audio uploads passing wait_timeout=600) stay cheap in practice while
|
|
still honoring the caller's intent if registration is unusually slow.
|
|
"""
|
|
test_file = tmp_path / "podcast.mp3"
|
|
test_file.write_bytes(b"fake audio")
|
|
|
|
# 3 RPCs: baseline + register + rename.
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None,
|
|
[[[["src_audio"]]]],
|
|
[[[["src_audio"], "Episode 1", [None, None, None, None, 10]]]],
|
|
]
|
|
|
|
sources_api._uploader.wait_until_registered = AsyncMock(
|
|
return_value=Source(id="src_audio", title="podcast.mp3", _type_code=10)
|
|
)
|
|
sources_api._uploader.wait_until_ready = AsyncMock()
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api.add_file(
|
|
"nb_123",
|
|
str(test_file),
|
|
title="Episode 1",
|
|
wait_timeout=600.0,
|
|
)
|
|
|
|
# wait_timeout is forwarded directly — no min() cap.
|
|
sources_api._uploader.wait_until_registered.assert_awaited_once_with(
|
|
"nb_123", "src_audio", timeout=600.0, transient_error_types=(10, 0, None)
|
|
)
|
|
sources_api._uploader.wait_until_ready.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_no_title_no_wait_does_not_wait(self, sources_api, mock_core, tmp_path):
|
|
"""Back-compat: wait=False with no custom title must NOT call
|
|
wait_until_ready — preserves the existing fast-return semantics.
|
|
"""
|
|
test_file = tmp_path / "test.pdf"
|
|
test_file.write_bytes(b"fake pdf content")
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_pdf"]]]]
|
|
sources_api._uploader.wait_until_ready = AsyncMock()
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file("nb_123", str(test_file))
|
|
|
|
assert result.id == "src_pdf"
|
|
assert result.title == "test.pdf"
|
|
sources_api._uploader.wait_until_ready.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_no_title_no_wait_returns_processing_status(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""The fire-and-forget placeholder Source returned when wait=False and
|
|
no custom title is supplied must carry status=PROCESSING. Previously
|
|
it defaulted to READY, so callers saw is_ready=True on a source that
|
|
had only just been registered. Regression guard for the CodeRabbit
|
|
comment on #396.
|
|
"""
|
|
test_file = tmp_path / "test.pdf"
|
|
test_file.write_bytes(b"fake pdf content")
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_pdf"]]]]
|
|
sources_api._uploader.wait_until_ready = AsyncMock()
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file("nb_123", str(test_file))
|
|
|
|
assert result.status == SourceStatus.PROCESSING
|
|
assert result.is_processing is True
|
|
assert result.is_ready is False
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("title", ["", " "])
|
|
async def test_add_file_rejects_blank_custom_title(
|
|
self, sources_api, mock_core, tmp_path, title
|
|
):
|
|
"""Blank custom titles should fail before upload or rename starts."""
|
|
test_file = tmp_path / "notes.md"
|
|
test_file.write_bytes(b"# notes\n")
|
|
|
|
with pytest.raises(ValidationError, match="Title cannot be empty"):
|
|
await sources_api.add_file("nb_123", str(test_file), title=title)
|
|
|
|
mock_core.rpc_executor.rpc_call.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_skips_rename_when_title_matches_filename(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""If the caller passes a title equal to the on-disk filename, no rename
|
|
RPC should be issued — saves a round-trip and avoids no-op writes.
|
|
"""
|
|
test_file = tmp_path / "report.pdf"
|
|
test_file.write_bytes(b"fake pdf content")
|
|
|
|
mock_core.rpc_executor.rpc_call.return_value = [[[["src_pdf"]]]]
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file("nb_123", str(test_file), title="report.pdf")
|
|
|
|
assert result.id == "src_pdf"
|
|
assert result.title == "report.pdf"
|
|
# No rename happened (title matches filename) — but registration is
|
|
# still 2 RPCs: baseline GET_NOTEBOOK + ADD_SOURCE_FILE.
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 2
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"rename_error",
|
|
[
|
|
RPCError("rename rpc blew up"),
|
|
NetworkError("rename network blew up"),
|
|
],
|
|
ids=["rpc", "network"],
|
|
)
|
|
async def test_add_file_rename_failure_keeps_upload(
|
|
self, sources_api, mock_core, tmp_path, caplog, rename_error
|
|
):
|
|
"""A failing rename must not lose the already-uploaded file; the
|
|
registered source is returned and a warning is logged.
|
|
"""
|
|
import logging
|
|
|
|
test_file = tmp_path / "doc.txt"
|
|
test_file.write_bytes(b"content")
|
|
|
|
# Registration succeeds; rename raises a library-level expected error
|
|
# (representative of what `self.rename` actually raises in the wild).
|
|
# The forced wait between register and rename is mocked separately.
|
|
# 3 RPCs: baseline + register + rename (the rename raises).
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None,
|
|
[[[["src_doc"]]]],
|
|
rename_error,
|
|
]
|
|
sources_api._uploader.wait_until_registered = AsyncMock(
|
|
return_value=Source(id="src_doc", title="doc.txt", _type_code=4)
|
|
)
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
with caplog.at_level(logging.WARNING, logger="notebooklm._sources"):
|
|
result = await sources_api.add_file("nb_123", str(test_file), title="Custom")
|
|
|
|
# The already-registered source is preserved.
|
|
assert result.id == "src_doc"
|
|
assert result.title == "doc.txt"
|
|
warning_records = [
|
|
rec for rec in caplog.records if "rename to 'Custom' failed" in rec.message
|
|
]
|
|
assert warning_records
|
|
assert warning_records[0].exc_info is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_with_title_preserves_waited_metadata(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""Renaming after the forced wait must not null out _type_code/url/created_at.
|
|
|
|
UPDATE_SOURCE's response shape can be sparse (see the rename() fallback
|
|
for ``result is None`` in _sources.py). If add_file naively overwrites
|
|
the fully-populated Source from wait_until_ready() with rename()'s
|
|
sparse return value, the source loses its real type code, URL, and
|
|
timestamp. Only the title should be taken from the rename response.
|
|
"""
|
|
from datetime import datetime, timezone
|
|
|
|
test_file = tmp_path / "podcast.mp3"
|
|
test_file.write_bytes(b"fake audio")
|
|
|
|
created_at = datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
# First rpc_call serves file registration. Second serves rename() —
|
|
# which returns a sparse Source (only id + new title) so we can verify
|
|
# the merge preserves type_code/url/created_at from the waited source.
|
|
# 3 RPCs: baseline + register + rename (rename returns None to
|
|
# trigger the fallback).
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None,
|
|
[[[["src_audio"]]]],
|
|
None, # Triggers rename()'s Source(id=source_id, title=new_title) fallback
|
|
]
|
|
sources_api._uploader.wait_until_registered = AsyncMock(
|
|
return_value=Source(
|
|
id="src_audio",
|
|
title="podcast.mp3",
|
|
_type_code=10,
|
|
url="https://example.com/audio",
|
|
created_at=created_at,
|
|
)
|
|
)
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file("nb_123", str(test_file), title="Episode 1")
|
|
|
|
# New title is applied...
|
|
assert result.title == "Episode 1"
|
|
# ...but the metadata populated by wait_until_ready() survives.
|
|
assert result._type_code == 10
|
|
assert result.url == "https://example.com/audio"
|
|
assert result.created_at == created_at
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_title_uses_narrow_wait_not_full_wait(
|
|
self, sources_api, mock_core, tmp_path
|
|
):
|
|
"""With wait=False but a custom title, add_file must call the narrow
|
|
wait_until_registered helper, NOT the full wait_until_ready. This is
|
|
the regression guard for the HIGH gemini-code-assist comment on #396:
|
|
callers that pass wait=False should not be blocked on full processing
|
|
just because they wanted to set a title.
|
|
"""
|
|
test_file = tmp_path / "long-audio.mp3"
|
|
test_file.write_bytes(b"fake audio")
|
|
|
|
# 3 RPCs: baseline + register + rename.
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None,
|
|
[[[["src_audio"]]]],
|
|
[[[["src_audio"], "My Title", [None, None, None, None, 10]]]],
|
|
]
|
|
|
|
sources_api._uploader.wait_until_registered = AsyncMock(
|
|
return_value=Source(id="src_audio", title="long-audio.mp3", _type_code=10)
|
|
)
|
|
sources_api._uploader.wait_until_ready = AsyncMock()
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
await sources_api.add_file(
|
|
"nb_123",
|
|
str(test_file),
|
|
title="My Title",
|
|
wait=False,
|
|
)
|
|
|
|
# Narrow wait was used...
|
|
sources_api._uploader.wait_until_registered.assert_awaited_once()
|
|
# ...and the full wait was NOT.
|
|
sources_api._uploader.wait_until_ready.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_file_rename_failure_still_waits(self, sources_api, mock_core, tmp_path):
|
|
"""A failed custom-title rename should not prevent wait=True polling.
|
|
With the new ordering, wait runs BEFORE rename, so this verifies the
|
|
wait still happens and the upload is preserved when the rename fails.
|
|
"""
|
|
test_file = tmp_path / "doc.txt"
|
|
test_file.write_bytes(b"content")
|
|
|
|
# 3 RPCs: baseline + register + rename (rename raises).
|
|
mock_core.rpc_executor.rpc_call.side_effect = [
|
|
None,
|
|
[[[["src_doc"]]]],
|
|
RPCError("rename rpc blew up"),
|
|
]
|
|
|
|
async def wait_side_effect(notebook_id, source_id, *, timeout, transient_error_types):
|
|
assert notebook_id == "nb_123"
|
|
assert source_id == "src_doc"
|
|
assert timeout == 120.0
|
|
assert transient_error_types == ()
|
|
# Wait runs BEFORE rename — baseline GET_NOTEBOOK + register
|
|
# RPCs have fired (2 total), no rename yet.
|
|
assert mock_core.rpc_executor.rpc_call.call_count == 2
|
|
return Source(id=source_id, title="doc.txt", _type_code=4)
|
|
|
|
sources_api._uploader.wait_until_ready = AsyncMock(side_effect=wait_side_effect)
|
|
|
|
mock_start_response = MagicMock()
|
|
mock_start_response.headers = {
|
|
"x-goog-upload-url": "https://notebooklm.google.com/upload/_/?upload_id=session"
|
|
}
|
|
mock_upload_response = MagicMock()
|
|
|
|
with patch("httpx.AsyncClient") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
mock_client.__aenter__.return_value = mock_client
|
|
mock_client.__aexit__.return_value = None
|
|
mock_client.post.side_effect = [mock_start_response, mock_upload_response]
|
|
mock_client_cls.return_value = mock_client
|
|
|
|
result = await sources_api.add_file(
|
|
"nb_123",
|
|
str(test_file),
|
|
title="Custom",
|
|
wait=True,
|
|
)
|
|
|
|
assert result.title == "doc.txt"
|
|
sources_api._uploader.wait_until_ready.assert_awaited_once_with(
|
|
"nb_123", "src_doc", timeout=120.0, transient_error_types=()
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# add_url() with YouTube detection tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestAddUrlWithYouTube:
|
|
"""Tests for add_url() with YouTube auto-detection."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_url_detects_youtube_and_uses_youtube_method(self, sources_api, mock_core):
|
|
"""Test that YouTube URLs are detected and routed correctly."""
|
|
mock_core.rpc_executor.rpc_call.return_value = [[["src_yt"], "YouTube Video"]]
|
|
|
|
await sources_api.add_url("nb_123", "https://youtu.be/dQw4w9WgXcQ")
|
|
|
|
# Check that the RPC was called with YouTube-specific params
|
|
call_args = mock_core.rpc_executor.rpc_call.call_args
|
|
params = call_args[0][1]
|
|
# YouTube params have the URL at position [0][0][7]
|
|
assert params[0][0][7] == ["https://youtu.be/dQw4w9WgXcQ"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_url_uses_regular_method_for_non_youtube(self, sources_api, mock_core):
|
|
"""Test that non-YouTube URLs use regular add method."""
|
|
mock_core.rpc_executor.rpc_call.return_value = [[["src_url"], "Example Site"]]
|
|
|
|
await sources_api.add_url("nb_123", "https://example.com/article")
|
|
|
|
# Check that the RPC was called with regular URL params
|
|
call_args = mock_core.rpc_executor.rpc_call.call_args
|
|
params = call_args[0][1]
|
|
# Regular URL params have the URL at position [0][0][2] (different from YouTube's [7])
|
|
assert params[0][0][2] == ["https://example.com/article"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_url_wraps_rpc_error_with_status_code(self, sources_api, mock_core):
|
|
"""When ADD_SOURCE returns null with a status code at wrb.fr[5] —
|
|
the #407 mode shared with #474 — the decoder raises ClientError /
|
|
RPCError and add_url wraps it into SourceAddError with the rich
|
|
diagnostic preserved as .cause. Previously allow_null=True on
|
|
_add_youtube_source swallowed the status code silently.
|
|
"""
|
|
from notebooklm.exceptions import ClientError, SourceAddError
|
|
|
|
mock_core.rpc_executor.rpc_call.side_effect = ClientError(
|
|
"RPC <id> returned null result with status code 7 (Permission denied). ...",
|
|
method_id="<id>",
|
|
rpc_code=7,
|
|
)
|
|
|
|
with pytest.raises(SourceAddError) as exc_info:
|
|
await sources_api.add_url("nb_123", "https://youtu.be/dQw4w9WgXcQ")
|
|
|
|
assert isinstance(exc_info.value.cause, ClientError)
|
|
assert exc_info.value.cause.rpc_code == 7
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"exc_factory",
|
|
[
|
|
lambda: __import__("notebooklm.exceptions", fromlist=["AuthError"]).AuthError(
|
|
"session expired"
|
|
),
|
|
lambda: __import__("notebooklm.exceptions", fromlist=["RateLimitError"]).RateLimitError(
|
|
"rate limited", method_id="<id>"
|
|
),
|
|
lambda: __import__("notebooklm.exceptions", fromlist=["ServerError"]).ServerError(
|
|
"backend down", method_id="<id>", rpc_code=503
|
|
),
|
|
],
|
|
ids=["auth", "rate-limit", "server"],
|
|
)
|
|
async def test_add_url_lets_transport_errors_propagate(
|
|
self, sources_api, mock_core, exc_factory
|
|
):
|
|
"""AuthError, RateLimitError, and ServerError must NOT be wrapped —
|
|
callers rely on the specific exception type for re-login / back-off
|
|
/ transient retry. Mirrors the propagation contract on
|
|
_register_file_source for #474.
|
|
"""
|
|
exc = exc_factory()
|
|
mock_core.rpc_executor.rpc_call.side_effect = exc
|
|
|
|
with pytest.raises(type(exc)):
|
|
await sources_api.add_url("nb_123", "https://example.com/article")
|
|
|
|
|
|
# =============================================================================
|
|
# _add_youtube_source() tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestAddYoutubeSource:
|
|
"""Tests for _add_youtube_source() helper."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_youtube_source_structure(self, sources_api, mock_core):
|
|
"""Test YouTube source params structure."""
|
|
mock_core.rpc_executor.rpc_call.return_value = [[["src_123"]]]
|
|
|
|
await sources_api._add_youtube_source("nb_123", "https://youtu.be/abc123")
|
|
|
|
call_args = mock_core.rpc_executor.rpc_call.call_args
|
|
params = call_args[0][1]
|
|
|
|
# Verify structure: [[None, None, None, ..., [url], None, None, 1]]
|
|
assert params[0][0][7] == ["https://youtu.be/abc123"]
|
|
assert params[0][0][10] == 1 # YouTube indicator
|
|
assert params[1] == "nb_123"
|
|
|
|
|
|
# =============================================================================
|
|
# _add_url_source() tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestAddUrlSource:
|
|
"""Tests for _add_url_source() helper."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_url_source_structure(self, sources_api, mock_core):
|
|
"""Test regular URL source params structure."""
|
|
mock_core.rpc_executor.rpc_call.return_value = [[["src_123"]]]
|
|
|
|
await sources_api._add_url_source("nb_123", "https://example.com/page")
|
|
|
|
call_args = mock_core.rpc_executor.rpc_call.call_args
|
|
params = call_args[0][1]
|
|
|
|
# Verify structure: URL at position 2 (different from YouTube which uses position 7)
|
|
assert params[0][0][2] == ["https://example.com/page"]
|
|
assert params[1] == "nb_123"
|
|
# Migrated to the nested trailing block (#1546): spec gains a trailing 1
|
|
# and [2],None,None collapses into [2,None,None,[1,...,[1]]].
|
|
assert params[0][0][-1] == 1
|
|
assert params[2] == [
|
|
2,
|
|
None,
|
|
None,
|
|
[1, None, None, None, None, None, None, None, None, None, [1]],
|
|
]
|
|
assert len(params) == 3
|