"""Tests for the NotebookLMClient periodic keepalive task.""" import asyncio import json import re import httpx import pytest from pytest_httpx import HTTPXMock import notebooklm._auth.keepalive as _auth_keepalive from notebooklm.auth import AuthTokens from notebooklm.client import NotebookLMClient from tests._helpers.client_factory import build_client_shell_for_tests ROTATE_URL_RE = re.compile(r"^https://accounts\.google\.com/RotateCookies$") @pytest.fixture def mock_auth(): """Create AuthTokens with no storage path (default tests don't persist).""" return AuthTokens( cookies={"SID": "test_sid", "__Secure-1PSIDTS": "test_1psidts", "HSID": "test_hsid"}, csrf_token="test_csrf", session_id="test_session", ) def _storage_auth(tmp_path) -> tuple[AuthTokens, "object"]: """Build AuthTokens backed by a real storage_state.json under tmp_path.""" storage_path = tmp_path / "storage_state.json" storage_path.write_text( json.dumps( { "cookies": [ { "name": "SID", "value": "initial_sid", "domain": ".google.com", "path": "/", }, { "name": "__Secure-1PSIDTS", "value": "test_1psidts", "domain": ".google.com", "path": "/", }, ] } ) ) auth = AuthTokens( cookies={"SID": "initial_sid", "__Secure-1PSIDTS": "test_1psidts"}, csrf_token="test_csrf", session_id="test_session", storage_path=storage_path, ) return auth, storage_path async def _read_storage_text_when_available(storage_path): """Read storage text, retrying transient Windows sharing violations.""" for _ in range(50): try: return storage_path.read_text() except PermissionError: await asyncio.sleep(0.05) return storage_path.read_text() async def _wait_until_storage_contains(storage_path, needle: str, failure_message: str) -> None: for _ in range(50): if needle in await _read_storage_text_when_available(storage_path): return await asyncio.sleep(0.05) pytest.fail(failure_message) def _rotate_requests(httpx_mock: HTTPXMock) -> list[httpx.Request]: return [r for r in httpx_mock.get_requests() if ROTATE_URL_RE.match(str(r.url))] async def _wait_for_rotate_requests( httpx_mock: HTTPXMock, *, minimum: int, failure_message: str, ) -> None: """Wait until the background keepalive task has emitted RotateCookies requests.""" requests: list[httpx.Request] = [] for _ in range(50): requests = _rotate_requests(httpx_mock) if len(requests) >= minimum: return await asyncio.sleep(0.05) requests = _rotate_requests(httpx_mock) if len(requests) < minimum: pytest.fail(f"{failure_message}; got {len(requests)}") class TestKeepaliveDisabledByDefault: @pytest.mark.asyncio async def test_keepalive_off_by_default(self, mock_auth, httpx_mock: HTTPXMock): """No keepalive task is spawned and no extra HTTP calls fire by default.""" client = NotebookLMClient(mock_auth) async with client: assert client._collaborators.lifecycle._keepalive_task is None # Give the loop a chance to run; nothing should happen await asyncio.sleep(0.1) # No RotateCookies request should have been issued for req in httpx_mock.get_requests(): assert "RotateCookies" not in str(req.url), f"Unexpected keepalive request: {req.url}" class TestKeepaliveLifecycle: @pytest.mark.asyncio @pytest.mark.no_default_keepalive_mock async def test_spawns_task_on_enter_cancels_on_exit(self, mock_auth, httpx_mock: HTTPXMock): """Task is created on __aenter__ and cleanly cancelled on __aexit__.""" httpx_mock.add_response( url=ROTATE_URL_RE, is_optional=True, is_reusable=True, status_code=204, ) client = NotebookLMClient( mock_auth, keepalive=0.05, keepalive_min_interval=0.01, ) async with client: task = client._collaborators.lifecycle._keepalive_task assert task is not None assert not task.done() # Task should be cleaned up; no warnings should be raised. assert client._collaborators.lifecycle._keepalive_task is None # Either cancelled or finished; never left dangling. assert task.done() class TestKeepaliveFloor: @pytest.mark.asyncio async def test_floor_clamps_low_interval(self, mock_auth): """Passing keepalive below keepalive_min_interval is clamped up.""" client = NotebookLMClient( mock_auth, keepalive=10.0, keepalive_min_interval=60.0, ) assert client._collaborators.lifecycle._keepalive_interval == 60.0 @pytest.mark.asyncio async def test_floor_does_not_lower_higher_interval(self, mock_auth): """Larger keepalive values pass through unchanged.""" client = NotebookLMClient( mock_auth, keepalive=600.0, keepalive_min_interval=60.0, ) assert client._collaborators.lifecycle._keepalive_interval == 600.0 @pytest.mark.asyncio async def test_none_keeps_disabled(self, mock_auth): """``keepalive=None`` keeps the loop disabled regardless of floor.""" client = NotebookLMClient( mock_auth, keepalive=None, keepalive_min_interval=60.0, ) assert client._collaborators.lifecycle._keepalive_interval is None class TestKeepaliveValidation: @pytest.mark.parametrize("bad", [0, 0.0, -1, -0.5, float("nan"), float("inf")]) def test_rejects_non_positive_or_non_finite_keepalive(self, mock_auth, bad): """``keepalive`` must be ``None`` or a positive finite number.""" with pytest.raises(ValueError, match="keepalive"): NotebookLMClient(mock_auth, keepalive=bad) @pytest.mark.parametrize("bad", [0, 0.0, -1, -0.5, float("nan"), float("inf")]) def test_rejects_non_positive_or_non_finite_floor(self, mock_auth, bad): """``keepalive_min_interval`` must be a positive finite number.""" with pytest.raises(ValueError, match="keepalive_min_interval"): NotebookLMClient(mock_auth, keepalive=120.0, keepalive_min_interval=bad) class TestKeepalivePokes: @pytest.mark.asyncio @pytest.mark.no_default_keepalive_mock async def test_pokes_at_interval(self, mock_auth, httpx_mock: HTTPXMock, monkeypatch): """At least two RotateCookies pokes fire within a short window. The test uses sub-second intervals to keep wall-clock cheap; the in-process rate-limit window (60 s in production) would otherwise suppress every iteration past the first. Patch the window down so the loop's pacing is the only thing being tested here. """ monkeypatch.setattr(_auth_keepalive, "_KEEPALIVE_RATE_LIMIT_SECONDS", 0.0) httpx_mock.add_response( url=ROTATE_URL_RE, is_optional=True, is_reusable=True, status_code=204, ) client = NotebookLMClient( mock_auth, keepalive=0.05, keepalive_min_interval=0.01, ) async with client: await _wait_for_rotate_requests( httpx_mock, minimum=2, failure_message="Expected at least 2 keepalive pokes", ) @pytest.mark.asyncio @pytest.mark.no_default_keepalive_mock async def test_failure_does_not_crash_loop(self, mock_auth, httpx_mock: HTTPXMock, monkeypatch): """A failing poke is swallowed and the loop continues. Same rate-limit-window patch as ``test_pokes_at_interval``: the sub-second test interval would otherwise be debounced into a single attempt by the in-process claim. """ monkeypatch.setattr(_auth_keepalive, "_KEEPALIVE_RATE_LIMIT_SECONDS", 0.0) # First poke: connection error. Subsequent pokes: 204. httpx_mock.add_exception( url=ROTATE_URL_RE, exception=httpx.ConnectError("simulated network blip"), ) httpx_mock.add_response( url=ROTATE_URL_RE, is_optional=True, is_reusable=True, status_code=204, ) client = NotebookLMClient( mock_auth, keepalive=0.05, keepalive_min_interval=0.01, ) async with client: await _wait_for_rotate_requests( httpx_mock, minimum=1, failure_message="Expected first keepalive poke", ) # Task is still running after the failure assert client._collaborators.lifecycle._keepalive_task is not None assert not client._collaborators.lifecycle._keepalive_task.done() await _wait_for_rotate_requests( httpx_mock, minimum=2, failure_message="Loop should have retried after failure", ) class TestKeepalivePersistenceFailure: @pytest.mark.asyncio @pytest.mark.no_default_keepalive_mock async def test_persistence_failure_logs_warning_and_continues( self, tmp_path, httpx_mock: HTTPXMock, monkeypatch, caplog ): """A failing ``save_cookies_to_storage`` is logged at WARNING; loop continues.""" auth, storage_path = _storage_auth(tmp_path) httpx_mock.add_response( url=ROTATE_URL_RE, is_optional=True, is_reusable=True, status_code=204, ) save_calls = [] def boom(cookies, path, **kwargs): save_calls.append(path) raise OSError("simulated disk full") # Phase 2 PR 4: inject the cookie-saver seam directly at the client # constructor rather than monkeypatching the legacy # ``notebooklm._core.save_cookies_to_storage`` indirection. client = NotebookLMClient( auth, keepalive=0.05, keepalive_min_interval=0.01, cookie_saver=boom, ) with caplog.at_level("WARNING", logger="notebooklm._core"): async with client: # Wait for at least 2 save attempts for _ in range(50): if len(save_calls) >= 2: break await asyncio.sleep(0.05) assert len(save_calls) >= 2, "Loop should retry after a persistence failure" warnings = [ r for r in caplog.records if r.levelname == "WARNING" and "persistence" in r.message ] assert warnings, "A persistence failure should surface as a WARNING log record" assert str(storage_path) in warnings[0].message class TestKeepaliveExplicitStoragePath: @pytest.mark.asyncio @pytest.mark.no_default_keepalive_mock async def test_explicit_storage_path_used_when_auth_lacks_one( self, tmp_path, httpx_mock: HTTPXMock ): """``NotebookLMClient(storage_path=...)`` is honored even when ``auth.storage_path`` is ``None``.""" # storage_state.json on disk, but auth.storage_path stays None storage_path = tmp_path / "storage_state.json" storage_path.write_text( json.dumps( { "cookies": [ { "name": "SID", "value": "manual_sid", "domain": ".google.com", "path": "/", }, ] } ) ) auth = AuthTokens( cookies={"SID": "manual_sid"}, csrf_token="t", session_id="s", # NOTE: storage_path is *not* set on auth ) httpx_mock.add_response( url=ROTATE_URL_RE, is_optional=True, is_reusable=True, status_code=204, headers={ "Set-Cookie": ( "__Secure-1PSIDTS=rotated_via_explicit; Domain=.google.com; Path=/; Secure" ), }, ) client = NotebookLMClient( auth, storage_path=storage_path, keepalive=0.05, keepalive_min_interval=0.01, ) async with client: await _wait_until_storage_contains( storage_path, "rotated_via_explicit", "Rotated cookie was not persisted to the explicit storage path", ) def test_explicit_storage_path_normalizes_onto_auth_without_mutating_caller(self, tmp_path): """The constructor exposes ``storage_path`` on ``client.auth`` so ``refresh_auth()`` and ``NotebookLMClient.close()`` (which read ``self._auth.storage_path`` directly, not the keepalive-specific path) persist to the same file. Crucially, the caller's original ``AuthTokens`` is *not* mutated, so reusing one ``AuthTokens`` across multiple ``NotebookLMClient`` instances with different storage paths is safe. """ storage_path = tmp_path / "storage_state.json" storage_path.write_text('{"cookies": []}') auth = AuthTokens( cookies={"SID": "x", "__Secure-1PSIDTS": "test_1psidts"}, csrf_token="t", session_id="s", # storage_path intentionally None ) assert auth.storage_path is None client = NotebookLMClient(auth, storage_path=storage_path) assert client.auth.storage_path == storage_path, ( "Explicit storage_path must be reflected on client.auth so non-keepalive " "code paths (refresh_auth, NotebookLMClient.close) see the same file" ) assert auth.storage_path is None, ( "Caller's AuthTokens must not be mutated — sharing one AuthTokens " "across clients with different storage paths must not leak between them" ) @pytest.mark.asyncio @pytest.mark.no_default_keepalive_mock async def test_close_persists_to_explicit_storage_path(self, tmp_path, httpx_mock: HTTPXMock): """``NotebookLMClient.close()`` calls ``save_cookies_to_storage`` with the explicit constructor ``storage_path`` even when keepalive never ran and ``auth.storage_path`` was ``None`` originally — proving the normalization actually wires the on-close save, not just the keepalive loop. """ storage_path = tmp_path / "storage_state.json" storage_path.write_text('{"cookies": []}') auth = AuthTokens( cookies={"SID": "x", "__Secure-1PSIDTS": "test_1psidts"}, csrf_token="t", session_id="s", # storage_path intentionally None ) save_calls: list[tuple[object, object]] = [] def spy(cookies, path, **kwargs): save_calls.append((cookies, path)) # Phase 2 PR 4: inject the cookie-saver seam directly. client = NotebookLMClient(auth, storage_path=storage_path, cookie_saver=spy) async with client: pass # no RPC calls; keepalive disabled by default # close()'s on-close save must have fired with the explicit storage_path assert any(call[1] == storage_path for call in save_calls), ( f"Expected close() to persist to {storage_path}, " f"but got: {[(type(c[0]).__name__, c[1]) for c in save_calls]}" ) class TestKeepalivePersistence: @pytest.mark.asyncio @pytest.mark.no_default_keepalive_mock async def test_persists_rotated_cookies_without_aexit(self, tmp_path, httpx_mock: HTTPXMock): """Rotated Set-Cookie from a poke is written to storage between pokes.""" auth, storage_path = _storage_auth(tmp_path) # The poke response sets a rotated 1PSIDTS on .google.com httpx_mock.add_response( url=ROTATE_URL_RE, is_optional=True, is_reusable=True, status_code=204, headers={ "Set-Cookie": ( "__Secure-1PSIDTS=rotated_value_xyz; Domain=.google.com; Path=/; Secure" ), }, ) client = NotebookLMClient( auth, keepalive=0.05, keepalive_min_interval=0.01, ) async with client: # Wait long enough for at least one poke + persist await _wait_until_storage_contains( storage_path, "rotated_value_xyz", "Rotated cookie was not persisted to storage_state.json before __aexit__ ran", ) # Sanity: the rotated cookie was written *while* the client was open, # not just at close time. data = json.loads(await _read_storage_text_when_available(storage_path)) cookie_names = {c["name"] for c in data["cookies"]} assert "__Secure-1PSIDTS" in cookie_names class TestSaveCookiesUnification: """Tests for NotebookLMClient.save_cookies — the single chokepoint that close, keepalive, and refresh_auth all route through.""" @pytest.mark.asyncio async def test_save_cookies_takes_in_process_lock_before_writing(self, tmp_path): """``NotebookLMClient.save_cookies`` holds ``_save_lock`` for the duration of the worker-thread write, so an older snapshot can't clobber a newer one within the same process.""" from notebooklm.client import NotebookLMClient auth = AuthTokens( cookies={"SID": "x", "__Secure-1PSIDTS": "test_1psidts"}, csrf_token="t", session_id="s", storage_path=tmp_path / "storage_state.json", ) (tmp_path / "storage_state.json").write_text('{"cookies": []}') lock_held_during_save: list[bool] = [] call_kwargs: list[dict] = [] core_ref: dict[str, NotebookLMClient] = {} def spy(jar, path, **kwargs): """Record lock state and the kwargs. Capturing ``kwargs`` here is the regression guard for the §3.4.1 fix: if ``NotebookLMClient.save_cookies`` ever stops threading ``original_snapshot=`` through, the assertion below catches it before production silently reverts to legacy merge. """ lock_held_during_save.append( core_ref["core"]._collaborators.cookie_persistence.save_lock.locked() ) call_kwargs.append(kwargs) return True # Phase 2 PR 4: inject the cookie-saver seam at construction. core = build_client_shell_for_tests(auth, cookie_saver=spy) core_ref["core"] = core await core._collaborators.lifecycle.save_cookies( core._collaborators.cookie_persistence, httpx.Cookies(), ) assert lock_held_during_save == [True], ( "save_cookies must hold _save_lock for the duration of " "save_cookies_to_storage so newer state always wins" ) assert len(call_kwargs) == 1 and "original_snapshot" in call_kwargs[0], ( "save_cookies must forward original_snapshot= to save_cookies_to_storage " "(dropping the kwarg would silently revert to legacy full-merge)" ) @pytest.mark.asyncio async def test_refresh_auth_routes_save_through_save_cookies( self, tmp_path, httpx_mock: HTTPXMock ): """``refresh_auth`` no longer calls ``save_cookies_to_storage`` directly; it routes through ``NotebookLMClient.save_cookies`` so the in-process lock is held — preventing an older keepalive snapshot from clobbering the freshly-refreshed tokens.""" storage_path = tmp_path / "storage_state.json" storage_path.write_text( json.dumps( { "cookies": [ { "name": "SID", "value": "x", "domain": ".google.com", "path": "/", }, { "name": "__Secure-1PSIDTS", "value": "test_1psidts", "domain": ".google.com", "path": "/", }, ] } ) ) auth = AuthTokens( cookies={"SID": "x", "__Secure-1PSIDTS": "test_1psidts", "HSID": "y"}, csrf_token="old_csrf", session_id="old_session", storage_path=storage_path, ) # NotebookLM homepage with new tokens (refresh_auth scrapes these) httpx_mock.add_response( url="https://notebooklm.google.com/", content=b'', ) save_calls: list[bool] = [] snapshot_kwarg_present: list[bool] = [] client_ref: dict[str, NotebookLMClient] = {} def spy(jar, path, **kwargs): """Record whether ``_save_lock`` is held when refresh_auth's save fires and whether ``original_snapshot=`` was forwarded. The snapshot-kwarg check is the §3.4.1 regression guard: refresh_auth must route through the snapshot/delta path, never the legacy full-merge path. """ save_calls.append( client_ref["client"]._collaborators.cookie_persistence.save_lock.locked() ) snapshot_kwarg_present.append("original_snapshot" in kwargs) return True # Phase 2 PR 4: inject the cookie-saver seam at construction. client = NotebookLMClient(auth, cookie_saver=spy) client_ref["client"] = client async with client: await client.refresh_auth() # At least one save (refresh_auth) plus close()'s on-close save. assert len(save_calls) >= 1 assert all(save_calls), ( "Every save fired during refresh_auth + close must be under the " f"in-process lock; got {save_calls}" ) assert all(snapshot_kwarg_present), ( "Every save fired during refresh_auth + close must forward " "original_snapshot= so the snapshot/delta path is taken, never " "the legacy full-merge fallback" ) class TestCrossProcessFileLock: """Tests for the OS-level file lock inside save_cookies_to_storage that serializes writes from different Python processes (e.g. an in-process keepalive plus a cron-driven `notebooklm auth refresh`).""" def test_save_cookies_to_storage_acquires_file_lock(self, tmp_path, monkeypatch): """``save_cookies_to_storage`` calls ``fcntl.flock(LOCK_EX)`` (POSIX) before reading or writing the storage file.""" import sys if sys.platform == "win32": pytest.skip("POSIX-specific test; Windows uses msvcrt.locking") import fcntl from notebooklm._auth.storage import save_cookies_to_storage storage_path = tmp_path / "storage_state.json" storage_path.write_text( '{"cookies": [{"name": "SID", "value": "old", "domain": ".google.com", "path": "/"}]}' ) flock_calls: list[int] = [] original_flock = fcntl.flock def spy_flock(fd, op): """Record each ``fcntl.flock`` operation while still performing it.""" flock_calls.append(op) return original_flock(fd, op) monkeypatch.setattr("fcntl.flock", spy_flock) from notebooklm._auth.storage import snapshot_cookie_jar jar = httpx.Cookies() empty_snapshot = snapshot_cookie_jar(jar) jar.set("SID", "new", domain=".google.com", path="/") save_cookies_to_storage(jar, storage_path, original_snapshot=empty_snapshot) assert fcntl.LOCK_EX in flock_calls, ( f"Expected an LOCK_EX call before the save, got: {flock_calls}" ) assert fcntl.LOCK_UN in flock_calls, ( f"Expected an LOCK_UN call after the save, got: {flock_calls}" ) def test_save_cookies_to_storage_creates_lock_sentinel(self, tmp_path): """The lock file is a sibling of the storage file with a `.lock` suffix so the storage file itself is free for the atomic temp-rename.""" from notebooklm._auth.storage import save_cookies_to_storage, snapshot_cookie_jar storage_path = tmp_path / "storage_state.json" storage_path.write_text( '{"cookies": [{"name": "SID", "value": "old", "domain": ".google.com", "path": "/"}]}' ) jar = httpx.Cookies() empty_snapshot = snapshot_cookie_jar(jar) jar.set("SID", "new", domain=".google.com", path="/") save_cookies_to_storage(jar, storage_path, original_snapshot=empty_snapshot) lock_path = storage_path.with_name(f".{storage_path.name}.lock") assert lock_path.exists(), f"Expected lock sentinel at {lock_path}"