/tests/unit/test_rate_limit_retry.py

"""429 retry budget on ``RpcExecutor.rpc_call``.

The rate-limit fix raises the ``rate_limit_max_retries`` default from
``0`` to ``3`` and adds capped exponential backoff as the sleep fallback
when 429 arrives without a parseable ``Retry-After`` header. Setting
``rate_limit_max_retries=0`` still restores raise-immediately behavior.
"""

from __future__ import annotations

import contextlib
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest

import notebooklm._deadline as _deadline
from notebooklm.rpc import RateLimitError, RPCError, RPCMethod
from tests._fixtures.kernel_test_helpers import install_http_client_for_test
from tests._helpers.client_factory import build_client_shell_for_tests
from tests.unit.conftest import install_post_as_stream


@pytest.fixture
def auth_tokens():
    auth = MagicMock()
    auth.csrf_token = "fake_csrf"
    return auth


def _build_429(retry_after: str | None = "1") -> MagicMock:
    resp = MagicMock(spec=httpx.Response)
    resp.status_code = 429
    resp.headers = {"retry-after": retry_after} if retry_after is not None else {}
    resp.reason_phrase = "Too Many Requests"

    def raise_429():
        raise httpx.HTTPStatusError("Rate Limit", request=MagicMock(), response=resp)

    resp.raise_for_status.side_effect = raise_429
    return resp


def _build_200(payload: list) -> MagicMock:
    resp = MagicMock(spec=httpx.Response)
    resp.status_code = 200
    resp.text = ")]}'\n[null,[" + str(payload).replace("'", '"') + "]]"
    resp.raise_for_status = MagicMock()
    return resp


@pytest.mark.asyncio
async def test_rate_limit_retry_success_with_budget(auth_tokens):
    """With budget>0 and a parseable Retry-After, the second call succeeds."""
    mock_client = AsyncMock(spec=httpx.AsyncClient)
    mock_client.post.side_effect = [_build_429("1"), _build_200([["result"]])]

    core = build_client_shell_for_tests(auth_tokens, rate_limit_max_retries=2)
    install_http_client_for_test(core._collaborators.kernel, mock_client)
    install_post_as_stream(None, mock_client, mock_client.post)

    # Decode may fail on the synthetic 200 — that's fine, what we care about
    # is the post counts and sleep budget. We expect either success or an
    # RPCError-tree decode failure, but the retry MUST have fired. Narrowed
    # from `except Exception` to keep unrelated programming errors visible.
    with patch("asyncio.sleep", AsyncMock()) as mock_sleep, contextlib.suppress(RPCError):
        await core._rpc_executor.rpc_call(RPCMethod.GET_NOTEBOOK, ["nb1"])

    assert mock_client.post.call_count == 2, (
        f"Expected initial 429 then 1 retry, got {mock_client.post.call_count}"
    )
    mock_sleep.assert_called_once_with(1)


@pytest.mark.asyncio
async def test_rate_limit_retry_after_larger_than_client_timeout_does_not_sleep(
    auth_tokens, monkeypatch
):
    """The middleware enforces the client timeout as the aggregate retry budget.

    A controlled monotonic clock starts the retry deadline at ``t=0`` and then
    jumps past the ``timeout`` budget before the budget check runs, so the
    middleware refuses to sleep (the requested retry would exceed the remaining
    budget) and re-raises ``RateLimitError`` after the single initial POST. The
    fake clock is *load-bearing*: with the real clock the near-zero elapsed time
    would leave the full budget intact and the middleware would sleep + retry.
    """
    mock_client = AsyncMock(spec=httpx.AsyncClient)
    # Retry-After (1s) is *within* the configured 10s budget, so the no-sleep
    # decision here is driven entirely by the controlled clock exhausting the
    # deadline — not by a header that trivially outruns the timeout.
    mock_client.post.return_value = _build_429("1")

    core = build_client_shell_for_tests(auth_tokens, timeout=10.0, rate_limit_max_retries=2)
    install_http_client_for_test(core._collaborators.kernel, mock_client)
    install_post_as_stream(None, mock_client, mock_client.post)

    sleeps: list[float] = []
    # A clock that leaps far past the 10s budget on *every* successive read, so
    # whichever read the retry deadline captures as its start, the next read
    # (the budget check) is already > timeout past it and ``remaining()``
    # collapses to 0. Robust to the deadline clock being shared with other
    # consumers in the call (the exact read ordering must not matter).
    clock = {"t": 0.0}

    def _monotonic() -> float:
        now = clock["t"]
        clock["t"] += 1000.0
        return now

    # ADR-0007: patch the deadline clock through a locally-imported module alias
    # (object form) rather than a ``notebooklm._deadline...`` string target.
    # ``MagicMock`` wraps the fake clock so the bite-check can assert the
    # injected seam was actually consulted. We replace ``_deadline.time`` with a
    # mock that wraps the real ``time`` module and overrides only ``monotonic``,
    # so the patch is isolated to ``_deadline``'s consumer-side binding and never
    # mutates the shared stdlib ``time`` module object for the rest of the process.
    fake_monotonic = MagicMock(side_effect=_monotonic)
    fake_time = MagicMock(wraps=_deadline.time)
    fake_time.monotonic = fake_monotonic
    monkeypatch.setattr(_deadline, "time", fake_time)

    async def _record_sleep(seconds: float) -> None:
        sleeps.append(seconds)

    with (
        patch("asyncio.sleep", side_effect=_record_sleep),
        pytest.raises(RateLimitError),
    ):
        await core._rpc_executor.rpc_call(RPCMethod.GET_NOTEBOOK, ["nb1"])

    # The controlled clock exhausted the budget before any retry sleep fired.
    assert mock_client.post.call_count == 1
    assert sleeps == []
    # Bite-check: the injected deadline clock was actually consulted.
    fake_monotonic.assert_called()


@pytest.mark.asyncio
async def test_rate_limit_retry_exhausted_with_budget(auth_tokens):
    """Budget=2 means: initial + 2 retries = 3 total posts before RateLimitError."""
    mock_client = AsyncMock(spec=httpx.AsyncClient)
    mock_client.post.return_value = _build_429("1")

    core = build_client_shell_for_tests(auth_tokens, rate_limit_max_retries=2)
    install_http_client_for_test(core._collaborators.kernel, mock_client)
    install_post_as_stream(None, mock_client, mock_client.post)

    with patch("asyncio.sleep", AsyncMock()) as mock_sleep, pytest.raises(RateLimitError):
        await core._rpc_executor.rpc_call(RPCMethod.GET_NOTEBOOK, ["nb1"])

    assert mock_client.post.call_count == 3
    assert mock_sleep.call_count == 2


@pytest.mark.asyncio
async def test_rate_limit_no_retry_if_disabled(auth_tokens):
    mock_client = AsyncMock(spec=httpx.AsyncClient)

    resp_429 = MagicMock(spec=httpx.Response)
    resp_429.status_code = 429
    resp_429.headers = {"retry-after": "1"}

    def raise_429():
        raise httpx.HTTPStatusError("Rate Limit", request=MagicMock(), response=resp_429)

    resp_429.raise_for_status.side_effect = raise_429

    mock_client.post.return_value = resp_429

    # Explicitly disable retries
    core = build_client_shell_for_tests(auth_tokens, rate_limit_max_retries=0)
    install_http_client_for_test(core._collaborators.kernel, mock_client)
    install_post_as_stream(None, mock_client, mock_client.post)

    with pytest.raises(RateLimitError):
        await core._rpc_executor.rpc_call(RPCMethod.GET_NOTEBOOK, ["nb1"])

    assert mock_client.post.call_count == 1


@pytest.mark.asyncio
async def test_rate_limit_exp_backoff_fallback_without_header(auth_tokens):
    """No Retry-After header → fall back to capped exponential backoff.

    Pre-fix, a 429 without ``Retry-After`` raised immediately even with
    budget>0. Audit §11 widened the retry circle: when the header is
    missing, the loop sleeps for ``min(2 ** attempt, 30)`` seconds with
    ±20% jitter and retries until the budget is exhausted, matching the
    5xx path.
    """
    mock_client = AsyncMock(spec=httpx.AsyncClient)
    mock_client.post.return_value = _build_429(retry_after=None)

    core = build_client_shell_for_tests(auth_tokens, rate_limit_max_retries=2)
    install_http_client_for_test(core._collaborators.kernel, mock_client)
    install_post_as_stream(None, mock_client, mock_client.post)

    sleeps: list[float] = []

    async def _record_sleep(seconds: float) -> None:
        sleeps.append(seconds)

    with patch("asyncio.sleep", side_effect=_record_sleep), pytest.raises(RateLimitError):
        await core._rpc_executor.rpc_call(RPCMethod.GET_NOTEBOOK, ["nb1"])

    # Initial + 2 retries = 3 POSTs before RateLimitError raises.
    assert mock_client.post.call_count == 3
    # 2 retries -> 2 backoff sleeps. Schedule: 2**0=1, 2**1=2, ±20% jitter.
    assert len(sleeps) == 2
    assert 0.8 <= sleeps[0] <= 1.2, f"attempt 1 backoff out of range: {sleeps[0]}"
    assert 1.6 <= sleeps[1] <= 2.4, f"attempt 2 backoff out of range: {sleeps[1]}"


@pytest.mark.asyncio
async def test_rate_limit_no_retry_without_header_when_disabled(auth_tokens):
    """``rate_limit_max_retries=0`` short-circuits even without a header."""
    mock_client = AsyncMock(spec=httpx.AsyncClient)
    mock_client.post.return_value = _build_429(retry_after=None)

    core = build_client_shell_for_tests(auth_tokens, rate_limit_max_retries=0)
    install_http_client_for_test(core._collaborators.kernel, mock_client)
    install_post_as_stream(None, mock_client, mock_client.post)

    with pytest.raises(RateLimitError):
        await core._rpc_executor.rpc_call(RPCMethod.GET_NOTEBOOK, ["nb1"])

    assert mock_client.post.call_count == 1


def test_rate_limit_max_retries_negative_raises(auth_tokens):
    """Negative budget is rejected at construction."""
    with pytest.raises(ValueError, match="rate_limit_max_retries must be >= 0"):
        build_client_shell_for_tests(auth_tokens, rate_limit_max_retries=-1)