learningcircuit--local-deep-research/tests/security/test_safe_requests_retries.py

"""Exponential-backoff retry wrapper for `safe_get`."""

import email.utils
from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock, patch

import pytest
import requests

from local_deep_research.security.safe_requests import (
    safe_get_with_retries,
)


def _mock_response(status_code=200, headers=None, content=b""):
    r = MagicMock(spec=requests.Response)
    r.status_code = status_code
    r.headers = headers or {}
    r.content = content
    return r


def _mock_response_body_raises(exc):
    """Response whose ``.content`` access raises ``exc`` on first read.

    Models the real production failure where ``safe_get`` returns a
    Response (headers received), then ``ChunkedEncodingError`` /
    ``ReadTimeout`` fires later when the body is consumed.
    """
    r = MagicMock(spec=requests.Response)
    r.status_code = 200
    r.headers = {}
    type(r).content = property(lambda self: (_ for _ in ()).throw(exc))
    return r


def test_first_attempt_success_returns_immediately():
    ok = _mock_response(200)
    with patch(
        "local_deep_research.security.safe_requests.safe_get",
        return_value=ok,
    ) as mock_get:
        resp = safe_get_with_retries("https://example.com/x")
    assert resp.status_code == 200
    assert mock_get.call_count == 1


def test_connection_error_retries_then_succeeds():
    ok = _mock_response(200)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[requests.ConnectionError("boom"), ok],
        ) as mock_get,
        patch(
            "local_deep_research.security.safe_requests.time.sleep"
        ) as mock_sleep,
    ):
        resp = safe_get_with_retries(
            "https://example.com/x", backoff_times=(0, 0, 0)
        )
    assert resp.status_code == 200
    assert mock_get.call_count == 2
    assert mock_sleep.called


def test_timeout_retries_then_gives_up():
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=requests.Timeout("slow"),
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        with pytest.raises(requests.Timeout):
            safe_get_with_retries(
                "https://example.com/x",
                max_retries=2,
                backoff_times=(0, 0, 0),
            )
    # 1 initial + 2 retries = 3 attempts total
    assert mock_get.call_count == 3


def test_http_500_retries_then_succeeds():
    bad = _mock_response(500)
    ok = _mock_response(200)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[bad, ok],
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        resp = safe_get_with_retries(
            "https://example.com/x", backoff_times=(0, 0, 0)
        )
    assert resp.status_code == 200
    assert mock_get.call_count == 2


def test_http_429_honors_retry_after():
    rate_limited = _mock_response(429, headers={"Retry-After": "7"})
    ok = _mock_response(200)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[rate_limited, ok],
        ),
        patch(
            "local_deep_research.security.safe_requests.time.sleep"
        ) as mock_sleep,
    ):
        safe_get_with_retries("https://example.com/x", backoff_times=(1, 1, 1))
    # Retry-After was 7, must override the 1-second schedule.
    assert mock_sleep.call_args[0][0] == 7


def test_ssrf_validation_error_is_not_retried():
    """ValueError means SSRF rejection — retrying would just re-fail."""
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=ValueError("URL failed security validation"),
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        with pytest.raises(ValueError):
            safe_get_with_retries(
                "http://169.254.169.254/", backoff_times=(0, 0, 0)
            )
    assert mock_get.call_count == 1


def test_http_404_is_not_retried():
    """4xx (other than 429) is caller's fault — no retry."""
    not_found = _mock_response(404)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            return_value=not_found,
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        resp = safe_get_with_retries(
            "https://example.com/x", backoff_times=(0, 0, 0)
        )
    assert resp.status_code == 404
    assert mock_get.call_count == 1


def test_retry_after_integer_is_capped_at_max():
    """A hostile Retry-After: 86400 must not pin the worker for a day."""
    rate_limited = _mock_response(429, headers={"Retry-After": "86400"})
    ok = _mock_response(200)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[rate_limited, ok],
        ),
        patch(
            "local_deep_research.security.safe_requests.time.sleep"
        ) as mock_sleep,
    ):
        safe_get_with_retries("https://example.com/x", backoff_times=(1, 1, 1))
    assert mock_sleep.call_args[0][0] == 300


def test_retry_after_http_date_form_is_parsed():
    """RFC 7231 HTTP-date form must be parsed, not silently ignored."""
    future = datetime.now(timezone.utc) + timedelta(seconds=60)
    http_date = email.utils.formatdate(future.timestamp(), usegmt=True)
    rate_limited = _mock_response(429, headers={"Retry-After": http_date})
    ok = _mock_response(200)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[rate_limited, ok],
        ),
        patch(
            "local_deep_research.security.safe_requests.time.sleep"
        ) as mock_sleep,
    ):
        safe_get_with_retries("https://example.com/x", backoff_times=(1, 1, 1))
    slept = mock_sleep.call_args[0][0]
    # 30s band absorbs clock jitter and test-harness latency.
    assert 45 <= slept <= 75, f"expected ~60s, got {slept}s"


def test_retry_after_unparseable_falls_back_to_schedule():
    """Garbage Retry-After falls back to the backoff schedule."""
    rate_limited = _mock_response(429, headers={"Retry-After": "garbage"})
    ok = _mock_response(200)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[rate_limited, ok],
        ),
        patch(
            "local_deep_research.security.safe_requests.time.sleep"
        ) as mock_sleep,
    ):
        safe_get_with_retries("https://example.com/x", backoff_times=(1, 2, 4))
    assert mock_sleep.call_args[0][0] == 1


def test_retry_after_negative_integer_is_clamped_to_zero():
    """time.sleep(-5) raises; Retry-After: -5 must be clamped to 0."""
    rate_limited = _mock_response(429, headers={"Retry-After": "-5"})
    ok = _mock_response(200)
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[rate_limited, ok],
        ),
        patch(
            "local_deep_research.security.safe_requests.time.sleep"
        ) as mock_sleep,
    ):
        safe_get_with_retries("https://example.com/x", backoff_times=(1, 1, 1))
    assert mock_sleep.call_args[0][0] == 0


# ---------------------------------------------------------------------------
# consume_body=True: retry on body-stream transients
# ---------------------------------------------------------------------------


def test_consume_body_default_does_not_read_body():
    """Default (``consume_body=False``) must not touch ``response.content``.

    The wrapper has historically returned the response without reading
    the body. Callers that stream large bodies (e.g., NDJSON line-by-line)
    rely on this — a regression that pre-reads would balloon memory.
    """
    body_marker = MagicMock()
    body_marker.__bool__ = MagicMock(side_effect=AssertionError("touched"))
    resp = _mock_response(200)
    type(resp).content = property(lambda self: body_marker)

    with patch(
        "local_deep_research.security.safe_requests.safe_get",
        return_value=resp,
    ):
        # Should not raise — default behavior leaves .content untouched.
        safe_get_with_retries("https://example.com/x")


def test_consume_body_retries_on_chunked_encoding_error():
    """Mid-stream ``ChunkedEncodingError`` retries the whole fetch."""
    bad = _mock_response_body_raises(
        requests.exceptions.ChunkedEncodingError(
            "IncompleteRead(835082 bytes read, 1262437 more expected)"
        )
    )
    ok = _mock_response(200, content=b"final")
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[bad, ok],
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        resp = safe_get_with_retries(
            "https://example.com/x",
            consume_body=True,
            backoff_times=(0, 0, 0),
        )
    assert resp.content == b"final"
    assert mock_get.call_count == 2
    bad.close.assert_called_once()


def test_consume_body_retries_on_read_timeout():
    """``ReadTimeout`` (Timeout but NOT ConnectionError) is also retried.

    Pinning this explicitly: the obvious ``except ConnectionError``
    catch would silently miss ReadTimeout (it's a Timeout subclass,
    not a ConnectionError subclass), so the implementation must list
    ``Timeout`` separately.
    """
    bad = _mock_response_body_raises(
        requests.exceptions.ReadTimeout("body read stalled")
    )
    ok = _mock_response(200, content=b"final")
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            side_effect=[bad, ok],
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        resp = safe_get_with_retries(
            "https://example.com/x",
            consume_body=True,
            backoff_times=(0, 0, 0),
        )
    assert resp.content == b"final"
    assert mock_get.call_count == 2


def test_consume_body_gives_up_after_max_retries():
    """Persistent body-read failures eventually surface to the caller."""
    bad = _mock_response_body_raises(
        requests.exceptions.ChunkedEncodingError("broken")
    )
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            return_value=bad,
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        with pytest.raises(requests.exceptions.ChunkedEncodingError):
            safe_get_with_retries(
                "https://example.com/x",
                consume_body=True,
                max_retries=2,
                backoff_times=(0, 0, 0),
            )
    # 1 initial + 2 retries = 3 attempts
    assert mock_get.call_count == 3


def test_consume_body_does_not_retry_value_error_from_body_guard():
    """``ValueError`` from oversized-body guard must NOT be retried.

    A 1 GB+ response isn't a transient network error — retrying just
    burns more bandwidth on the same outcome. The guard's ValueError
    must propagate immediately on the first attempt.
    """
    bad = _mock_response_body_raises(
        ValueError("Response body too large: >1073741825 bytes")
    )
    with (
        patch(
            "local_deep_research.security.safe_requests.safe_get",
            return_value=bad,
        ) as mock_get,
        patch("local_deep_research.security.safe_requests.time.sleep"),
    ):
        with pytest.raises(ValueError):
            safe_get_with_retries(
                "https://example.com/x",
                consume_body=True,
                backoff_times=(0, 0, 0),
            )
    # Single attempt — no retry on ValueError.
    assert mock_get.call_count == 1


def test_consume_body_returns_cached_body_to_caller():
    """After a successful retry, ``response.content`` is cached.

    The point of consuming inside the loop is so the caller doesn't
    repeat the read (and risk a second transient). Verify the body
    is available without a second read.
    """
    ok = _mock_response(200, content=b"hello world")
    with patch(
        "local_deep_research.security.safe_requests.safe_get",
        return_value=ok,
    ):
        resp = safe_get_with_retries("https://example.com/x", consume_body=True)
    assert resp.content == b"hello world"