teng-lin--notebooklm-py/tests/unit/test_profile_atomic_write.py

"""Unit tests for unified atomic profile-state write (P1-20).

The pre-P1-20 layout stored Google auth state across two files:

* ``storage_state.json`` — Playwright cookie/origin state.
* ``context.json`` — sibling JSON with ``{"account": {"authuser", "email"}}``
  and CLI notebook/conversation context.

Each file was atomically written on its own. Between the two writes there was
a window in which an external reader (e.g. a long-running ``notebooklm chat``
in a sibling shell) could see either the new cookies bundled with the old
account metadata, or vice versa. Tier-9 P1-20 closes this by bundling the
``account`` record INTO ``storage_state.json`` under a ``notebooklm``
namespace key:

    {
      "cookies": [...],
      "origins": [...],
      "notebooklm": {"version": 1, "account": {"authuser": 1, "email": "..."}}
    }

A single ``atomic_write_json`` is now the only commit point for the
(cookies, account) pair. ``context.json`` keeps holding non-account CLI
state (``notebook_id``, ``conversation_id``); the account key, if still
present from legacy installs, is migrated lazily on next write.

This module covers:

1. Round-trip of the new unified record.
2. Migration: a legacy two-file fixture reads cleanly under the new reader.
3. Migration: writing account metadata after a legacy read drops the
   ``account`` key from ``context.json`` (preserving other CLI state).
4. Torn-write fault injection: if ``storage_state.json`` write fails after
   cookies + account were serialized into the same temp file, the original
   on-disk file is preserved untouched (no half-written state).
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pytest

from notebooklm._auth.account import (
    clear_account_metadata,
    get_account_email_for_storage,
    get_authuser_for_storage,
    read_account_metadata,
    write_account_metadata,
)


def _write_storage_state(path: Path, payload: dict[str, Any]) -> None:
    path.write_text(json.dumps(payload), encoding="utf-8")


def _read_storage_state(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _context_path(storage_path: Path) -> Path:
    return storage_path.with_name("context.json")


def test_write_account_metadata_lands_inline_in_storage_state(tmp_path: Path) -> None:
    """The post-P1-20 writer puts account inside storage_state.json."""
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(
        storage_path,
        {
            "cookies": [{"name": "SID", "value": "v", "domain": ".google.com", "path": "/"}],
            "origins": [],
        },
    )
    write_account_metadata(storage_path, authuser=1, email="alice@example.com")
    payload = _read_storage_state(storage_path)
    assert payload["notebooklm"]["version"] == 1
    assert payload["notebooklm"]["account"] == {"authuser": 1, "email": "alice@example.com"}


def test_read_account_metadata_prefers_in_band_record(tmp_path: Path) -> None:
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(
        storage_path,
        {
            "cookies": [],
            "origins": [],
            "notebooklm": {
                "version": 1,
                "account": {"authuser": 2, "email": "bob@example.com"},
            },
        },
    )
    assert get_authuser_for_storage(storage_path) == 2
    assert get_account_email_for_storage(storage_path) == "bob@example.com"


def test_legacy_two_file_fixture_reads_cleanly(tmp_path: Path) -> None:
    """ACCEPTANCE-CRITICAL: existing two-file profile loads under new reader.

    If this migration test breaks for any existing user, they will lose their
    account binding (authuser/email) on the next CLI run — Tier-9 P1-20 PR
    body must surface this.
    """
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(
        storage_path,
        {
            "cookies": [{"name": "SID", "value": "v", "domain": ".google.com", "path": "/"}],
            "origins": [],
        },
    )
    # Legacy sibling context.json with account but no in-band record.
    _context_path(storage_path).write_text(
        json.dumps(
            {
                "account": {"authuser": 3, "email": "charlie@example.com"},
                "notebook_id": "nb-123",
            }
        ),
        encoding="utf-8",
    )

    assert get_authuser_for_storage(storage_path) == 3
    assert get_account_email_for_storage(storage_path) == "charlie@example.com"
    assert read_account_metadata(storage_path) == {
        "authuser": 3,
        "email": "charlie@example.com",
    }


def test_migration_on_write_removes_legacy_account_key_only(tmp_path: Path) -> None:
    """After upgrade write, ``context.json`` keeps non-account state.

    ``context.json`` also holds ``notebook_id`` / ``conversation_id`` —
    migration must NOT clobber those.
    """
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(storage_path, {"cookies": [], "origins": []})
    _context_path(storage_path).write_text(
        json.dumps(
            {
                "account": {"authuser": 4, "email": "dana@example.com"},
                "notebook_id": "nb-456",
                "conversation_id": "conv-789",
            }
        ),
        encoding="utf-8",
    )

    # Trigger a unified write; this should migrate the legacy record.
    write_account_metadata(storage_path, authuser=5, email="erin@example.com")

    in_band = _read_storage_state(storage_path)["notebooklm"]["account"]
    assert in_band == {"authuser": 5, "email": "erin@example.com"}

    # Non-account context state preserved.
    legacy = json.loads(_context_path(storage_path).read_text(encoding="utf-8"))
    assert "account" not in legacy
    assert legacy.get("notebook_id") == "nb-456"
    assert legacy.get("conversation_id") == "conv-789"


def test_in_band_account_overrides_legacy_account(tmp_path: Path) -> None:
    """When both forms exist, in-band wins."""
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(
        storage_path,
        {
            "cookies": [],
            "origins": [],
            "notebooklm": {
                "version": 1,
                "account": {"authuser": 7, "email": "new@example.com"},
            },
        },
    )
    _context_path(storage_path).write_text(
        json.dumps({"account": {"authuser": 1, "email": "stale@example.com"}}),
        encoding="utf-8",
    )

    assert get_authuser_for_storage(storage_path) == 7
    assert get_account_email_for_storage(storage_path) == "new@example.com"


def test_clear_account_metadata_clears_in_band(tmp_path: Path) -> None:
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(
        storage_path,
        {
            "cookies": [],
            "origins": [],
            "notebooklm": {
                "version": 1,
                "account": {"authuser": 9, "email": "zoe@example.com"},
            },
        },
    )
    clear_account_metadata(storage_path)
    # Either the namespace is gone, or its account is gone — either way the
    # reader reports no account. The file itself remains valid JSON.
    _read_storage_state(storage_path)  # sanity-check the file still parses
    assert read_account_metadata(storage_path) == {}
    assert get_authuser_for_storage(storage_path) == 0
    assert get_account_email_for_storage(storage_path) is None


def test_clear_account_metadata_clears_legacy_two_file(tmp_path: Path) -> None:
    """Backward compat: clearing still removes legacy ``context.json`` account."""
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(storage_path, {"cookies": [], "origins": []})
    _context_path(storage_path).write_text(
        json.dumps(
            {
                "account": {"authuser": 8, "email": "leah@example.com"},
                "notebook_id": "nb-keep",
            }
        ),
        encoding="utf-8",
    )
    clear_account_metadata(storage_path)
    # Non-account context preserved.
    legacy = json.loads(_context_path(storage_path).read_text(encoding="utf-8"))
    assert "account" not in legacy
    assert legacy.get("notebook_id") == "nb-keep"


def test_torn_write_fault_injection_preserves_original(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """ACCEPTANCE-CRITICAL: writer crash between cookies+metadata is recoverable.

    If we crash during the unified write, the original on-disk
    ``storage_state.json`` must remain valid and unchanged — no half-written
    state mixing new cookies with old account or vice versa.
    """
    storage_path = tmp_path / "storage_state.json"
    original_payload = {
        "cookies": [{"name": "SID", "value": "old-v", "domain": ".google.com", "path": "/"}],
        "origins": [],
        "notebooklm": {
            "version": 1,
            "account": {"authuser": 1, "email": "original@example.com"},
        },
    }
    _write_storage_state(storage_path, original_payload)

    # Inject a crash mid-write: the next ``os.replace`` raises.
    import notebooklm._atomic_io as atomic_io_mod

    original_replace = atomic_io_mod.os.replace

    def _boom(src: Any, dst: Any) -> None:
        raise OSError("simulated crash during atomic replace")

    monkeypatch.setattr(atomic_io_mod.os, "replace", _boom)

    with pytest.raises(OSError, match="simulated crash"):
        write_account_metadata(storage_path, authuser=99, email="never@example.com")

    monkeypatch.setattr(atomic_io_mod.os, "replace", original_replace)

    # Original file untouched: reader sees the OLD record consistently.
    assert _read_storage_state(storage_path) == original_payload
    assert get_authuser_for_storage(storage_path) == 1
    assert get_account_email_for_storage(storage_path) == "original@example.com"

    # No torn-write temp files leaked beside the storage file. The filelock
    # sentinel (``.storage_state.json.lock``) is expected to persist — filelock
    # >= 3.29 no longer unlinks it on release — and is not a torn-write leak.
    leftover_temps = [
        p for p in tmp_path.glob(".storage_state.json.*") if p.name != ".storage_state.json.lock"
    ]
    assert leftover_temps == [], f"Temp file leaked: {leftover_temps}"


def test_torn_write_reader_never_sees_mixed_state(tmp_path: Path) -> None:
    """Property: at any observation point, reader sees old XOR new — never mixed.

    Atomicity comes from ``atomic_write_json``'s single ``os.replace``: until
    the rename commits, the reader sees only the previous on-disk version.
    This test asserts the contract is exercised by the unified write path —
    after a successful write, both cookies and account come from the same
    commit; if the write fails, neither rolls forward.
    """
    storage_path = tmp_path / "storage_state.json"
    # Round 1: write old state.
    write_account_metadata(storage_path, authuser=10, email="round1@example.com")
    # Sanity: the file now exists with the round-1 record.
    payload_1 = _read_storage_state(storage_path)
    assert payload_1["notebooklm"]["account"]["email"] == "round1@example.com"

    # Round 2: overwrite with a new account record.
    write_account_metadata(storage_path, authuser=20, email="round2@example.com")
    payload_2 = _read_storage_state(storage_path)
    # The reader sees exactly one record — never a merge.
    assert payload_2["notebooklm"]["account"] == {
        "authuser": 20,
        "email": "round2@example.com",
    }
    # The non-account file structure round-trips unchanged.
    assert payload_2.get("cookies") == payload_1.get("cookies")


def test_unified_format_version_is_pinned_to_one(tmp_path: Path) -> None:
    """Pin the version number so any future bump is intentional."""
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(storage_path, {"cookies": [], "origins": []})
    write_account_metadata(storage_path, authuser=1, email="v@example.com")
    payload = _read_storage_state(storage_path)
    assert payload["notebooklm"]["version"] == 1


def test_write_without_email_omits_email_field(tmp_path: Path) -> None:
    """Default-account login: authuser=0, no email — record omits email."""
    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(storage_path, {"cookies": [], "origins": []})
    write_account_metadata(storage_path, authuser=0, email=None)
    payload = _read_storage_state(storage_path)
    assert payload["notebooklm"]["account"] == {"authuser": 0}


def test_write_preserves_cookies_and_origins(tmp_path: Path) -> None:
    """Writing account metadata MUST NOT touch cookies / origins."""
    storage_path = tmp_path / "storage_state.json"
    cookies = [
        {"name": "SID", "value": "v1", "domain": ".google.com", "path": "/"},
        {"name": "HSID", "value": "v2", "domain": ".google.com", "path": "/"},
    ]
    origins = [
        {"origin": "https://notebooklm.google.com", "localStorage": [{"name": "k", "value": "v"}]}
    ]
    _write_storage_state(storage_path, {"cookies": cookies, "origins": origins})
    write_account_metadata(storage_path, authuser=1, email="alice@example.com")
    payload = _read_storage_state(storage_path)
    assert payload["cookies"] == cookies
    assert payload["origins"] == origins


def _capture_account_lock_path(monkeypatch: pytest.MonkeyPatch) -> dict[str, Path]:
    """Record the ``storage_state.json`` lock path account.py passes to ``FileLock``.

    The captured contextmanager is a no-op so the surrounding read-modify-write
    still completes against the real file. ``write_account_metadata`` also calls
    ``_drop_legacy_account_key``, which locks the sibling ``context.json.lock``;
    that path is filtered out so it can't clobber the value we assert on. Paths
    are canonicalized (``.expanduser().resolve()``) because they back
    cross-process locking, where distinct spellings of the same file must
    compare equal.
    """
    import contextlib

    from notebooklm._auth import account as account_mod

    seen: dict[str, Path] = {}

    @contextlib.contextmanager
    def _fake_filelock(path: str, *args: Any, **kwargs: Any):  # type: ignore[no-untyped-def]
        resolved = Path(path).expanduser().resolve()
        if "storage_state.json" in resolved.name:
            seen["path"] = resolved
        yield

    monkeypatch.setattr(account_mod, "FileLock", _fake_filelock)
    return seen


def _capture_storage_lock_path(monkeypatch: pytest.MonkeyPatch) -> dict[str, Path]:
    """Record the lock path storage.py passes to ``_file_lock_exclusive``.

    Paths are canonicalized to match ``_capture_account_lock_path`` so the two
    captures compare equal even if their spellings differ.
    """
    import contextlib

    from notebooklm._auth import storage as storage_mod

    seen: dict[str, Path] = {}

    @contextlib.contextmanager
    def _fake_exclusive(lock_path: Path):  # type: ignore[no-untyped-def]
        seen["path"] = Path(lock_path).expanduser().resolve()
        yield

    monkeypatch.setattr(storage_mod, "_file_lock_exclusive", _fake_exclusive)
    return seen


def test_storage_state_mutators_share_one_lock_file(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """All ``storage_state.json`` mutators must serialize on the SAME flock file.

    Tier-0 data-loss regression: ``save_cookies_to_storage`` (storage.py) used
    the dotted ``.storage_state.json.lock`` sibling while ``account.py``'s
    metadata writers used the non-dotted ``storage_state.json.lock`` — a
    DIFFERENT file. Cookie-save and account-metadata writes therefore locked on
    distinct files and could lose updates under concurrency. Assert both paths
    derive an identical lock filename for the same ``storage_state.json``.
    """
    import httpx

    from notebooklm._auth.storage import save_cookies_to_storage

    storage_path = tmp_path / "storage_state.json"
    _write_storage_state(storage_path, {"cookies": [], "origins": []})

    # account.py: write_account_metadata
    account_seen = _capture_account_lock_path(monkeypatch)
    write_account_metadata(storage_path, authuser=1, email="alice@example.com")
    account_write_lock = account_seen["path"]

    # account.py: _clear_in_band_account (via clear_account_metadata)
    account_seen.clear()
    clear_account_metadata(storage_path)
    account_clear_lock = account_seen["path"]

    # storage.py: save_cookies_to_storage (the canonical cookie-save writer)
    storage_seen = _capture_storage_lock_path(monkeypatch)
    save_cookies_to_storage(
        httpx.Cookies(),
        path=storage_path,
        original_snapshot={},
    )
    storage_cookie_lock = storage_seen["path"]

    # Canonical name is the dotted, hidden sibling (storage.py contract).
    # Resolve to match the canonicalized captures above.
    expected = storage_path.with_name(f".{storage_path.name}.lock").expanduser().resolve()
    assert storage_cookie_lock == expected
    assert account_write_lock == expected
    assert account_clear_lock == expected
    # And, transitively, all three agree on the exact same file on disk.
    assert account_write_lock == account_clear_lock == storage_cookie_lock