"""Unit tests for the saved-from-chat CREATE_NOTE encoder (issue #660). These tests exercise ``_chat.notes.build_save_chat_as_note_params`` and pin its output against the wire-captured payload at ``tests/unit/fixtures/save_chat_as_note_create_note_request.json``. The golden test is the most important one: a byte-exact match (via deep-equal of the decoded JSON structure) guarantees that the encoder produces the same payload Google's web UI sends when its "Save to note" button is clicked. Drift from that payload risks the server silently dropping citation anchors and reverting the note to plain text. The encoder moved from ``_mind_map.py`` to ``_chat/notes.py`` in Phase 6 (refactor-history.md Step 8, ADR-0013); the test imports were updated accordingly. """ from __future__ import annotations import json from pathlib import Path import pytest from notebooklm._chat.notes import ( _CITATION_MARKER_RE, _resolve_reference, _strip_citation_markers, build_save_chat_as_note_params, ) from notebooklm.types import ChatReference FIXTURES_DIR = Path(__file__).parent / "fixtures" def _load_request_fixture() -> dict: return json.loads((FIXTURES_DIR / "save_chat_as_note_create_note_request.json").read_text()) class TestStripCitationMarkers: """Tests for the ``[N]`` marker stripping helper.""" def test_single_marker_with_leading_space(self): clean, positions = _strip_citation_markers("One fruit mentioned is apples [1].") assert clean == "One fruit mentioned is apples." # Marker [1] was at "apples"+space+marker position; after stripping, # the position in the clean text where the marker was is just after # "apples" = char 29 (exclusive end of "One fruit mentioned is apples"). assert positions == [(1, 29)] def test_no_markers(self): clean, positions = _strip_citation_markers("plain answer, no citations") assert clean == "plain answer, no citations" assert positions == [] def test_multiple_markers_segmented(self): clean, positions = _strip_citation_markers("X is true [1]. Y is also true [2].") # Stripping " [1]" then " [2]" yields "X is true. Y is also true." assert clean == "X is true. Y is also true." # [1] anchors clean[0..9] = "X is true"; [2] anchors clean[0..25] = the whole # text before the trailing period assert positions == [(1, 9), (2, 25)] def test_adjacent_markers(self): clean, positions = _strip_citation_markers("X [1][2].") # First match consumes " [1]" with leading space, leaving "X[2]." # Second match consumes "[2]" with NO leading space (regex space is optional) assert clean == "X." assert positions == [(1, 1), (2, 1)] def test_marker_at_start(self): clean, positions = _strip_citation_markers("[1] starts here.") # No leading space; the optional " " doesn't match assert clean == " starts here." assert positions == [(1, 0)] def test_regex_compiled(self): # Sanity: ensure the module-level compile didn't drift. assert _CITATION_MARKER_RE.pattern == r" ?\[(\d+)\]" class TestBuildSaveChatAsNoteParamsGolden: """Byte-exact golden tests against the captured curl payload. The fixture was decoded from the user-supplied curl request sent by the NotebookLM web UI when "Save to note" was clicked on a chat answer with a single citation. The encoder MUST match it exactly: if our output diverges in shape, ordering, or scalar types (e.g. emitting ``false`` where the wire uses ``0``), the server may silently strip citation anchors. """ def test_golden_single_citation(self): fixture = _load_request_fixture() expected_params = fixture["params"] notebook_id = expected_params[0] answer_text = expected_params[1] title = expected_params[4] # Reconstruct the ChatReference the web UI would have built from # the same chat response. Slot [3][0] of the captured request # carries the source-passage descriptor; we read source_id / # chunk_id / span / cited_text out of it. passage = expected_params[3][0] cited_text = passage[4][0][0][2][0][0][2][0] # follow the nested wrapper # passage[3] = [[null, start, end]]; passage[5] = [[[passage_id], source_id]]; # passage[6] = [chunk_id] start_char = passage[3][0][1] end_char = passage[3][0][2] passage_id = passage[5][0][0][0] source_id = passage[5][0][1] chunk_id = passage[6][0] references = [ ChatReference( source_id=source_id, citation_number=1, cited_text=cited_text, start_char=start_char, end_char=end_char, chunk_id=chunk_id, passage_id=passage_id, ) ] actual_params = build_save_chat_as_note_params(notebook_id, answer_text, references, title) # Deep-equal of the structured form. We also check the compact # JSON serialization to catch boolean / null divergence that # deep-equal would miss (Python's `False == 0` quirk). assert actual_params == expected_params # Serialize both with the same separators the rpc encoder uses # (rpc/encoder.py:39: separators=(",", ":")) and compare strings. # This is the strict byte-exact check. actual_json = json.dumps(actual_params, separators=(",", ":")) expected_json = json.dumps(expected_params, separators=(",", ":")) assert actual_json == expected_json def test_passage_id_falls_back_to_chunk_id_when_unset(self): """When ChatReference.passage_id is None (the production path, since the chat parser doesn't currently surface this UUID), the encoder fills the 4th-UUID slot with chunk_id as a placeholder.""" ref = ChatReference( source_id="src-1", citation_number=1, cited_text="text", start_char=0, end_char=4, chunk_id="chunk-1", passage_id=None, ) params = build_save_chat_as_note_params("nb-id", "X [1].", [ref], "Title") # Slot [3][0][5][0] = [[passage_id], source_id] # With passage_id=None, the encoder uses chunk_id. passage_descriptor = params[3][0] assert passage_descriptor[5][0][0] == ["chunk-1"] assert passage_descriptor[5][0][1] == "src-1" # The keyed version in rich_content uses the same placeholder. assert params[5][3][0][1][5][0][0] == ["chunk-1"] def test_passage_id_overrides_chunk_id_when_set(self): """When ChatReference.passage_id IS set (forward-compat path), the encoder uses it for the 4th-UUID slot regardless of chunk_id.""" ref = ChatReference( source_id="src-1", citation_number=1, cited_text="text", start_char=0, end_char=4, chunk_id="chunk-1", passage_id="passage-real-uuid", ) params = build_save_chat_as_note_params("nb-id", "X [1].", [ref], "Title") passage_descriptor = params[3][0] assert passage_descriptor[5][0][0] == ["passage-real-uuid"] assert passage_descriptor[5][0][1] == "src-1" def test_encoder_serializes_booleans_as_zero_not_false(self): """Regression guard: every "rendering flags" slot the ENCODER emits must serialize to ``0`` (integer), never ``false`` (boolean). Python's ``json.dumps(False)`` emits ``false`` while ``json.dumps(0)`` emits ``0`` — the wire payload uses ``0``, and the server's strict request channel won't normalize this. We assert on the encoder's actual output (not on the static fixture) so the check can't pass tautologically: if a future change makes any slot emit ``True``/``False``, this test catches it.""" ref = ChatReference( source_id="src-1", citation_number=1, cited_text="passage text", start_char=0, end_char=12, chunk_id="chunk-1", ) params = build_save_chat_as_note_params( "nb-id", "Answer with citation [1].", [ref], "Title" ) actual_json = json.dumps(params, separators=(",", ":")) assert "false" not in actual_json assert "true" not in actual_json assert "[0,0,0,null,null,null,null,0,0]" in actual_json class TestBuildSaveChatAsNoteParamsBehavior: """Behavioural tests beyond the single-citation golden.""" def _make_ref(self, n: int, chunk_id: str, source_id: str) -> ChatReference: return ChatReference( source_id=source_id, citation_number=n, cited_text=f"passage text {n}", start_char=0, end_char=14, chunk_id=chunk_id, ) def test_empty_references_raises(self): with pytest.raises(ValueError, match="non-empty references"): build_save_chat_as_note_params("nb-id", "no citations here.", [], "Title") def test_references_without_chunk_id_raises(self): ref = ChatReference(source_id="src-1", chunk_id=None) with pytest.raises(ValueError, match="chunk_id"): build_save_chat_as_note_params("nb-id", "X [1].", [ref], "Title") def test_multi_citation_produces_one_anchor_per_marker(self): refs = [ self._make_ref(1, "chunk-a", "src-a"), self._make_ref(2, "chunk-b", "src-b"), ] params = build_save_chat_as_note_params( "nb-id", "First fact [1]. Second fact [2].", refs, "Title" ) # Two unique chunks → two source_passages entries assert len(params[3]) == 2 # Two markers → two chunk_refs in rich_content[0][1] chunk_refs = params[5][0][1] assert len(chunk_refs) == 2 # Cumulative spans: [1] anchors clean[0..10] = "First fact" # [2] anchors clean[0..23] = "First fact. Second fact" (23 chars) assert chunk_refs[0] == [["chunk-a"], [None, 0, 10]] assert chunk_refs[1] == [["chunk-b"], [None, 0, 23]] def test_dedupes_source_passages_by_chunk_id(self): # Same chunk_id under two distinct citation numbers (the realistic # repeat: one chunk cited at two markers) → only one source_passages # entry, but each [N] marker still gets its own anchor. (Previously # this pinned the numbered-positional fallback — both refs numbered 1 # and [2] resolved via references[1]; that fallback is now gated to # unnumbered candidates so a numbering hole can't mis-anchor.) refs = [ self._make_ref(1, "chunk-shared", "src-shared"), self._make_ref(2, "chunk-shared", "src-shared"), ] params = build_save_chat_as_note_params("nb-id", "A [1] B [2]", refs, "Title") assert len(params[3]) == 1 # deduped # Two markers found in answer_text → two chunk_refs assert len(params[5][0][1]) == 2 def test_marker_without_matching_reference_is_skipped(self, caplog): # Answer has [1] and [99] but references only covers [1] refs = [self._make_ref(1, "chunk-1", "src-1")] with caplog.at_level("WARNING"): params = build_save_chat_as_note_params("nb-id", "A [1] B [99]", refs, "Title") chunk_refs = params[5][0][1] # Only one anchor — the [99] marker is silently dropped with a warning assert len(chunk_refs) == 1 assert any("[99]" in record.message or "99" in record.message for record in caplog.records) def test_holed_numbering_skips_missing_marker_instead_of_mis_anchoring(self, caplog): """Regression (codex review of the citation hardening): no wrong anchors. The wire parser keeps RAW ordinals when it skips a malformed citation row, so ``[good#1, skipped#2, good#3]`` reaches the encoder as refs numbered ``{1, 3}``. Marker ``[2]`` must resolve to nothing — the old positional fallback would have anchored ``references[1]`` = raw #3, a silently WRONG anchor. Missing anchor (with a warning) is the only acceptable outcome. """ refs = [ self._make_ref(1, "chunk-1", "src-1"), self._make_ref(3, "chunk-3", "src-3"), ] # Resolver level: exact matches work, the hole yields None. resolved_1 = _resolve_reference(refs, 1) assert resolved_1 is not None and resolved_1.chunk_id == "chunk-1" assert _resolve_reference(refs, 2) is None resolved_3 = _resolve_reference(refs, 3) assert resolved_3 is not None and resolved_3.chunk_id == "chunk-3" # Encoder level: [2]'s anchor is dropped with a warning; [1] and [3] # anchor their own chunks. with caplog.at_level("WARNING"): params = build_save_chat_as_note_params("nb-id", "A [1] B [2] C [3]", refs, "Title") chunk_refs = params[5][0][1] assert [anchor[0] for anchor in chunk_refs] == [["chunk-1"], ["chunk-3"]] assert any("[2]" in record.message for record in caplog.records) def test_positional_fallback_still_serves_unnumbered_references(self): """The legacy unset-number contract is preserved: positional lookup applies when the positional candidate carries no citation_number.""" refs = [ ChatReference(source_id="src-a", chunk_id="chunk-a"), ChatReference(source_id="src-b", chunk_id="chunk-b"), ] resolved = _resolve_reference(refs, 2) assert resolved is not None and resolved.chunk_id == "chunk-b" def test_seven_element_params_shape(self): """Sanity: top-level params is always exactly 7 elements.""" refs = [self._make_ref(1, "c1", "s1")] params = build_save_chat_as_note_params("nb-id", "X [1].", refs, "T") assert len(params) == 7 assert params[2] == [2] # mode flag assert params[6] == [2] # trailer flag assert params[0] == "nb-id" assert params[1] == "X [1]." assert params[4] == "T" def test_rich_content_has_five_slots(self): refs = [self._make_ref(1, "c1", "s1")] params = build_save_chat_as_note_params("nb-id", "X [1].", refs, "T") rich = params[5] assert len(rich) == 5 assert rich[1] is None # always-null slot assert rich[2] is None # always-null slot assert rich[4] == 1 # trailer flag def test_source_span_and_local_text_wrapper_diverge_for_nonzero_start(self): """Regression: source-document span (slot [3]) uses ref.start_char / ref.end_char (absolute offsets in the source), but the text-passage wrapper (slot [4]) uses LOCAL offsets [0, len(cited_text)]. The captured single-citation fixture has start_char=0 and end_char == len(cited_text), so they coincidentally match — but real refs from the chat parser commonly have non-zero source offsets (e.g. chars 100..200 in a long document). If the encoder used ref.end_char for both, the local wrapper would emit ``[[0, 200, ...]]`` for a 100-char passage, breaking server-side hover anchoring.""" ref = ChatReference( source_id="src-1", citation_number=1, cited_text="exactly-ten", # 11 chars, deliberately != source span start_char=100, end_char=111, # source offsets chunk_id="chunk-1", ) params = build_save_chat_as_note_params("nb-id", "X [1].", [ref], "Title") passage = params[3][0] # Source span uses absolute offsets. assert passage[3] == [[None, 100, 111]] # Text wrapper uses LOCAL offsets — [0, len(cited_text)]. passage_group = passage[4][0] assert passage_group[0][0] == 0 assert passage_group[0][1] == 11 # len("exactly-ten") # The deeply-nested inner mirror must also use the local end. inner_offsets = passage_group[0][2][0][0] assert inner_offsets[0] == 0 assert inner_offsets[1] == 11 def test_empty_cited_text_collapses_span(self): """Regression: when ref.cited_text is empty, span must collapse to [0, 0] rather than emitting an invalid [None, start, 0] where start > 0 (e.g. start_char=10 + cited_text='').""" ref = ChatReference( source_id="src-1", citation_number=1, cited_text="", start_char=10, end_char=20, chunk_id="chunk-1", ) params = build_save_chat_as_note_params("nb-id", "X [1].", [ref], "Title") passage = params[3][0] # Source span collapses to [0, 0], not [10, 0]. assert passage[3] == [[None, 0, 0]] # Local text wrapper end is also 0. assert passage[4][0][0][1] == 0 def test_title_passthrough(self): refs = [self._make_ref(1, "c1", "s1")] params = build_save_chat_as_note_params("nb-id", "X [1].", refs, "My Custom Title") assert params[4] == "My Custom Title"