teng-lin--notebooklm-py/tests/integration/test_auto_refresh.py

"""Integration tests for automatic token refresh."""

import asyncio
from unittest.mock import MagicMock

import httpx
import pytest

from notebooklm import NotebookLMClient
from notebooklm.auth import AuthTokens
from notebooklm.rpc import RPCError, RPCMethod
from tests.integration.conftest import install_post_as_stream

# mock-based refresh-callback wiring tests; no HTTP, no cassette.
# Opt out of the tier-enforcement hook in tests/integration/conftest.py.
pytestmark = pytest.mark.allow_no_vcr


class TestAutoRefreshIntegration:
    @pytest.mark.asyncio
    async def test_client_has_refresh_callback_wired(self):
        """NotebookLMClient should wire refresh_auth as callback."""
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        # Bound methods aren't identical, so compare underlying function
        assert client._collaborators.auth_coord._refresh_callback is not None
        assert (
            client._collaborators.auth_coord._refresh_callback.__func__
            is NotebookLMClient.refresh_auth
        )
        # ``_refresh_lock`` is lazily created on first ``_await_refresh``.
        # At construction time it is ``None`` so the client can be
        # instantiated outside a running loop; the helper allocates the
        # lock on demand inside the async refresh path.
        assert client._collaborators.auth_coord._refresh_lock is None

    @pytest.mark.asyncio
    async def test_full_refresh_flow_http_error(self):
        """Test complete auto-refresh flow for HTTP 401 errors."""
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="old_csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        # Override retry delay for faster tests
        client._composed.chain_host._refresh_retry_delay = 0

        # Track refresh calls
        refresh_calls = []

        async def tracking_refresh():
            refresh_calls.append(True)
            # Simulate successful refresh
            client._auth.csrf_token = "new_csrf"
            # Wave 3 of plan ``host-protocol-removal`` deleted the
            # Session-level ``update_auth_headers`` forward; call the
            # canonical coordinator method directly with explicit kwargs.
            client._collaborators.auth_coord.update_auth_headers(
                auth=client._auth,
                kernel=client._collaborators.kernel,
            )
            return client._auth

        client._collaborators.auth_coord._refresh_callback = tracking_refresh

        # Mock HTTP responses
        call_count = [0]

        async def mock_post(*args, **kwargs):
            call_count[0] += 1
            if call_count[0] == 1:
                # First call: simulate HTTP 401
                request = httpx.Request("POST", args[0])
                response = httpx.Response(401, request=request)
                raise httpx.HTTPStatusError("Unauthorized", request=request, response=response)
            # Second call: success
            response = MagicMock()
            response.text = ')]}\'\\n[["wrb.fr","wXbhsf",[[[["nb1"],["Notebook 1"]]]]]]'
            response.raise_for_status = MagicMock()
            return response

        # Override the runtime decode-response seam before the RPC fires.
        client._seams.decode_response = lambda *a, **kw: [[["nb1"], ["Notebook 1"]]]

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)
            await client.notebooks.list()

        assert len(refresh_calls) == 1, "Should have refreshed once"
        assert call_count[0] == 2, "Should have retried once"

    @pytest.mark.asyncio
    async def test_full_refresh_flow_rpc_error(self):
        """Test complete auto-refresh flow for RPC auth errors."""
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="old_csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        client._composed.chain_host._refresh_retry_delay = 0

        refresh_calls = []

        async def tracking_refresh():
            refresh_calls.append(True)
            client._auth.csrf_token = "new_csrf"
            # Wave 3 of plan ``host-protocol-removal`` deleted the
            # Session-level ``update_auth_headers`` forward; call the
            # canonical coordinator method directly with explicit kwargs.
            client._collaborators.auth_coord.update_auth_headers(
                auth=client._auth,
                kernel=client._collaborators.kernel,
            )
            return client._auth

        client._collaborators.auth_coord._refresh_callback = tracking_refresh

        # Mock HTTP to succeed, but decode_response to fail with auth error first
        async def mock_post(*args, **kwargs):
            response = MagicMock()
            response.text = "mock response"
            response.raise_for_status = MagicMock()
            return response

        decode_count = [0]

        def mock_decode(*args, **kwargs):
            decode_count[0] += 1
            if decode_count[0] == 1:
                raise RPCError("Authentication expired")
            return [[["nb1"], ["Notebook 1"]]]

        # Override the runtime decode-response seam before the RPC fires.
        client._seams.decode_response = mock_decode

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)
            await client.notebooks.list()

        assert len(refresh_calls) == 1, "Should have refreshed once"
        assert decode_count[0] == 2, "Should have retried once"

    @pytest.mark.asyncio
    async def test_wire_401_then_decoded_auth_error_refreshes_once(self):
        """Issue #1205: a wire-401 followed by a decoded auth error on the SAME
        logical call must drive exactly ONE refresh.

        Before consolidation the HTTP-status layer (``AuthRefreshMiddleware``)
        and the decoded-RPC layer (``RpcExecutor``) tracked their once-per-call
        guard independently — the chain's per-request ``auth_refreshed`` flag
        and the executor's ``_is_retry`` flag could not see each other. So a
        ``401 → refresh#1 → 200 → decoded-auth-error → refresh#2`` sequence
        refreshed twice. The shared :class:`RefreshBudget` threaded through both
        layers now bounds the logical call to a single refresh; the decoded
        auth error surfaces to the caller instead of triggering a second
        refresh.
        """
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="old_csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        client._composed.chain_host._refresh_retry_delay = 0

        refresh_calls = []

        async def tracking_refresh():
            refresh_calls.append(True)
            client._auth.csrf_token = "new_csrf"
            client._collaborators.auth_coord.update_auth_headers(
                auth=client._auth,
                kernel=client._collaborators.kernel,
            )
            return client._auth

        client._collaborators.auth_coord._refresh_callback = tracking_refresh

        call_count = [0]

        async def mock_post(*args, **kwargs):
            call_count[0] += 1
            if call_count[0] == 1:
                # Wire 401 → HTTP-status layer refreshes (refresh #1) and
                # retries the POST.
                request = httpx.Request("POST", args[0])
                response = httpx.Response(401, request=request)
                raise httpx.HTTPStatusError("Unauthorized", request=request, response=response)
            # The post-refresh retry returns HTTP 200; the decoded payload
            # still carries an auth error.
            response = MagicMock()
            response.text = "mock response"
            response.raise_for_status = MagicMock()
            return response

        auth_rpc_error = RPCError("authentication expired")

        def mock_decode(*args, **kwargs):
            raise auth_rpc_error

        client._seams.decode_response = mock_decode

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)

            # The decoded auth error surfaces — the shared budget was already
            # spent by the HTTP-status refresh, so the decoded layer does NOT
            # refresh again and re-raises the original auth error.
            with pytest.raises(RPCError) as raised:
                await client.notebooks.list()

        assert raised.value is auth_rpc_error
        assert len(refresh_calls) == 1, "wire-401 + decoded-auth-error must refresh exactly once"
        # Two POSTs: the initial 401 and the single post-refresh retry. No
        # third POST, because the decoded layer did not refresh-and-retry.
        assert call_count[0] == 2

    @pytest.mark.asyncio
    async def test_decoded_auth_retry_increments_auth_retry_metric(self):
        """Issue #1205: the decoded-RPC refresh layer now counts the auth retry.

        Before consolidation only the HTTP-status layer incremented
        ``rpc_auth_retries``; the decode-time refresh-and-retry leg silently
        skipped it. The shared refresh body counts on both layers.
        """
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="old_csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        client._composed.chain_host._refresh_retry_delay = 0

        async def tracking_refresh():
            client._auth.csrf_token = "new_csrf"
            client._collaborators.auth_coord.update_auth_headers(
                auth=client._auth,
                kernel=client._collaborators.kernel,
            )
            return client._auth

        client._collaborators.auth_coord._refresh_callback = tracking_refresh

        async def mock_post(*args, **kwargs):
            response = MagicMock()
            response.text = "mock response"
            response.raise_for_status = MagicMock()
            return response

        decode_count = [0]

        def mock_decode(*args, **kwargs):
            decode_count[0] += 1
            if decode_count[0] == 1:
                raise RPCError("Authentication expired")
            return [[["nb1"], ["Notebook 1"]]]

        client._seams.decode_response = mock_decode

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)
            await client.notebooks.list()

        assert client._collaborators.metrics.snapshot().rpc_auth_retries == 1

    @pytest.mark.asyncio
    async def test_refresh_delay_is_applied(self):
        """Test that retry delay is actually applied."""
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        client._composed.chain_host._refresh_retry_delay = 0.1  # 100ms delay

        async def mock_refresh():
            return auth

        client._collaborators.auth_coord._refresh_callback = mock_refresh

        call_count = [0]

        async def mock_post(*args, **kwargs):
            call_count[0] += 1
            if call_count[0] == 1:
                request = httpx.Request("POST", args[0])
                response = httpx.Response(401, request=request)
                raise httpx.HTTPStatusError("Unauthorized", request=request, response=response)
            response = MagicMock()
            response.text = "mock"
            response.raise_for_status = MagicMock()
            return response

        # Override the runtime decode-response seam before the RPC fires.
        client._seams.decode_response = lambda *a, **kw: []

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)

            start_time = asyncio.get_event_loop().time()
            await client.notebooks.list()
            elapsed = asyncio.get_event_loop().time() - start_time

        # Should have taken at least the delay time
        assert elapsed >= 0.09, f"Delay should be applied, elapsed: {elapsed}"

    @pytest.mark.asyncio
    async def test_no_retry_on_cookie_expiration(self):
        """Test that full cookie expiration is not retried (requires re-login)."""
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        client._composed.chain_host._refresh_retry_delay = 0

        async def failing_refresh():
            # Simulates refresh_auth detecting redirect to login
            raise ValueError("Authentication expired. Run 'notebooklm login' to re-authenticate.")

        client._collaborators.auth_coord._refresh_callback = failing_refresh

        async def mock_post(*args, **kwargs):
            request = httpx.Request("POST", args[0])
            response = httpx.Response(401, request=request)
            raise httpx.HTTPStatusError("Unauthorized", request=request, response=response)

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)

            # Should raise the original HTTP error with refresh failure as cause
            with pytest.raises(httpx.HTTPStatusError) as exc_info:
                await client.notebooks.list()

            assert exc_info.value.__cause__ is not None
            assert "re-authenticate" in str(exc_info.value.__cause__)

    @pytest.mark.asyncio
    async def test_http_auth_error_does_not_replay_non_idempotent_write(self):
        """A mid-flight 401 on a non-idempotent create is NOT replayed.

        Regression for issue #1157. ``CREATE_NOTEBOOK`` is PROBE_THEN_CREATE,
        so ``resolve_effective_disable_internal_retries`` forces the effective
        disable flag True. The server may have committed the notebook before
        the 401 surfaced, so ``AuthRefreshMiddleware`` must NOT refresh and
        re-POST — that would duplicate the notebook. The original auth error
        propagates so ``NotebooksAPI.create``'s probe-then-create wrapper can
        disambiguate. Driven through the public ``client.notebooks.create``
        surface so the regression is pinned end-to-end.
        """
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="old_csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        client._composed.chain_host._refresh_retry_delay = 0

        refresh_calls = []

        async def tracking_refresh():
            refresh_calls.append(True)
            return client._auth

        client._collaborators.auth_coord._refresh_callback = tracking_refresh

        create_post_count = [0]

        async def mock_post(*args, **kwargs):
            url = args[0]
            # ``NotebooksAPI.create`` first lists notebooks to capture a
            # baseline; that LIST_NOTEBOOKS POST must succeed so only the
            # CREATE_NOTEBOOK leg exercises the auth-error path.
            if RPCMethod.LIST_NOTEBOOKS.value in str(url):
                response = MagicMock()
                response.text = "list-ok"
                response.raise_for_status = MagicMock()
                return response
            create_post_count[0] += 1
            request = httpx.Request("POST", url)
            response = httpx.Response(401, request=request)
            raise httpx.HTTPStatusError("Unauthorized", request=request, response=response)

        # Baseline ``list()`` decodes to an empty notebook list; the create's
        # decode never runs because the POST raises a 401 first.
        client._seams.decode_response = lambda *a, **kw: []

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)

            with pytest.raises(RPCError):
                await client.notebooks.create("My Notebook")

        assert refresh_calls == [], "non-idempotent write must not trigger an auth refresh"
        assert create_post_count[0] == 1, "CREATE_NOTEBOOK must POST exactly once (no replay)"

    @pytest.mark.asyncio
    async def test_rpc_auth_error_does_not_replay_non_idempotent_write(self):
        """A decoded auth-shaped ``RPCError`` is NOT replayed for a create.

        Regression for issue #1157 — the decode-time refresh-and-retry leg in
        ``RpcExecutor`` must honor the effective disable classification just
        like the HTTP-status leg. ``CREATE_NOTEBOOK`` resolves to disabled
        retries, so the decoded auth error surfaces without a second POST.
        Driven through the public ``client.notebooks.create`` surface.
        """
        auth = AuthTokens(
            cookies={"SID": "test"},
            csrf_token="old_csrf",
            session_id="sid",
        )

        client = NotebookLMClient(auth)
        client._composed.chain_host._refresh_retry_delay = 0

        refresh_calls = []

        async def tracking_refresh():
            refresh_calls.append(True)
            return client._auth

        client._collaborators.auth_coord._refresh_callback = tracking_refresh

        create_post_count = [0]

        async def mock_post(*args, **kwargs):
            if RPCMethod.CREATE_NOTEBOOK.value in str(args[0]):
                create_post_count[0] += 1
            response = MagicMock()
            response.text = "mock response"
            response.raise_for_status = MagicMock()
            return response

        create_decode_count = [0]

        def mock_decode(raw, rpc_id, *args, **kwargs):
            # The baseline ``list()`` decodes to an empty list; the create's
            # decode raises an auth-shaped RPCError to exercise the
            # decode-time refresh-and-retry leg.
            if rpc_id == RPCMethod.CREATE_NOTEBOOK.value:
                create_decode_count[0] += 1
                raise RPCError("Authentication expired")
            return []

        client._seams.decode_response = mock_decode

        async with client:
            install_post_as_stream(None, client._collaborators.kernel.get_http_client(), mock_post)

            with pytest.raises(RPCError):
                await client.notebooks.create("My Notebook")

        assert refresh_calls == [], "non-idempotent write must not trigger an auth refresh"
        assert create_post_count[0] == 1, "CREATE_NOTEBOOK must POST exactly once (no replay)"
        assert create_decode_count[0] == 1, (
            "CREATE_NOTEBOOK decode must run once — no retried decode"
        )