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

"""Tests for _check_response_size() and SafeSession.send() intermediate redirect branching.

Covers:
- _check_response_size() direct unit tests (8 tests) — zero prior coverage
- SafeSession.send() intermediate redirect size-check logic (10 tests) — lines 437-443
"""

import pytest
from unittest.mock import patch, MagicMock

import requests

from local_deep_research.security import ssrf_validator
from local_deep_research.security.safe_requests import (
    _check_response_size,
    SafeSession,
    MAX_RESPONSE_SIZE,
    _REDIRECT_STATUS_CODES,
)


def _mock_response(status_code=200, headers=None):
    """Build a MagicMock Response with the given status and headers."""
    resp = MagicMock(spec=requests.Response)
    resp.status_code = status_code
    resp.headers = headers or {}
    resp.url = "https://example.com"
    # _check_response_size may install a body guard via _install_body_guard,
    # which accesses response.raw.read
    resp.raw = MagicMock()
    return resp


# ── _check_response_size() direct unit tests ────────────────────────


class TestCheckResponseSizeDirect:
    """Direct unit tests for _check_response_size (lines 29-53)."""

    def test_no_content_length_header(self):
        """No Content-Length header → no raise."""
        resp = _mock_response(headers={})
        _check_response_size(resp)  # should not raise

    def test_empty_string_content_length(self):
        """Empty string Content-Length is falsy → no raise."""
        resp = _mock_response(headers={"Content-Length": ""})
        _check_response_size(resp)

    def test_negative_content_length(self):
        """Negative Content-Length → returns silently (line 46-47)."""
        resp = _mock_response(headers={"Content-Length": "-1"})
        _check_response_size(resp)

    def test_non_numeric_content_length(self):
        """Non-numeric Content-Length → ValueError caught, returns (line 44)."""
        resp = _mock_response(headers={"Content-Length": "abc"})
        _check_response_size(resp)

    def test_content_length_type_error(self):
        """Content-Length value that causes TypeError in int() → caught (line 44)."""
        resp = _mock_response(headers={"Content-Length": None})
        _check_response_size(resp)

    def test_exactly_at_max_size(self):
        """Content-Length exactly at MAX_RESPONSE_SIZE → no raise (> not >=)."""
        resp = _mock_response(
            headers={"Content-Length": str(MAX_RESPONSE_SIZE)}
        )
        _check_response_size(resp)

    def test_one_over_max_raises(self):
        """Content-Length one over MAX → raises ValueError, response.close() called."""
        resp = _mock_response(
            headers={"Content-Length": str(MAX_RESPONSE_SIZE + 1)}
        )
        with pytest.raises(ValueError, match="Response too large"):
            _check_response_size(resp)
        resp.close.assert_called_once()

    def test_normal_size_under_limit(self):
        """Normal Content-Length well under limit → no raise."""
        resp = _mock_response(headers={"Content-Length": "1024"})
        _check_response_size(resp)


# ── SafeSession.send() intermediate redirect branching ──────────────


def _prep_request(url="https://example.com"):
    """Build a minimal PreparedRequest."""
    prep = requests.PreparedRequest()
    prep.prepare_url(url, {})
    prep.prepare_method("GET")
    return prep


@pytest.fixture
def _bypass_ssrf():
    """Patch validate_url to always return True (bypass SSRF check)."""
    with patch.object(
        ssrf_validator,
        "validate_url",
        return_value=True,
    ):
        yield


class TestSendIntermediateRedirectBranching:
    """Tests for is_intermediate_redirect logic (lines 437-443)."""

    def test_302_with_location_and_allow_redirects_applies_size_check(
        self, _bypass_ssrf
    ):
        """302 + Location + allow_redirects=True → size check still applies."""
        oversized = _mock_response(
            302,
            {
                "Location": "https://other.com",
                "Content-Length": str(MAX_RESPONSE_SIZE + 1),
            },
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request(), allow_redirects=True)

    def test_302_with_location_and_allow_redirects_false_applies_size_check(
        self, _bypass_ssrf
    ):
        """302 + Location + allow_redirects=False → apply size check → raises."""
        oversized = _mock_response(
            302,
            {
                "Location": "https://other.com",
                "Content-Length": str(MAX_RESPONSE_SIZE + 1),
            },
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request(), allow_redirects=False)

    def test_302_without_location_applies_size_check(self, _bypass_ssrf):
        """302 without Location header → apply size check → raises."""
        oversized = _mock_response(
            302, {"Content-Length": str(MAX_RESPONSE_SIZE + 1)}
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request(), allow_redirects=True)

    def test_302_without_location_normal_size_passes(self, _bypass_ssrf):
        """302 without Location, normal Content-Length → no raise."""
        resp = _mock_response(302, {"Content-Length": "512"})
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=resp):
            result = session.send(_prep_request(), allow_redirects=True)
            assert result.status_code == 302

    def test_200_normal_size_passes(self, _bypass_ssrf):
        """200 response, normal size → passes."""
        resp = _mock_response(200, {"Content-Length": "1024"})
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=resp):
            result = session.send(_prep_request())
            assert result.status_code == 200

    def test_200_oversized_raises(self, _bypass_ssrf):
        """200 response, oversized Content-Length → raises."""
        oversized = _mock_response(
            200, {"Content-Length": str(MAX_RESPONSE_SIZE + 1)}
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request())

    @pytest.mark.parametrize("status_code", [301, 302, 303, 307, 308])
    def test_all_redirect_codes_apply_size_check(
        self, _bypass_ssrf, status_code
    ):
        """All 5 redirect codes + Location + oversized → rejected by size check."""
        oversized = _mock_response(
            status_code,
            {
                "Location": "https://other.com",
                "Content-Length": str(MAX_RESPONSE_SIZE + 1),
            },
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request(), allow_redirects=True)

    def test_allow_redirects_defaults_to_true_still_checks_size(
        self, _bypass_ssrf
    ):
        """No allow_redirects in kwargs → size check still applies."""
        oversized = _mock_response(
            302,
            {
                "Location": "https://other.com",
                "Content-Length": str(MAX_RESPONSE_SIZE + 1),
            },
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request())

    def test_304_not_in_redirect_set_applies_size_check(self, _bypass_ssrf):
        """304 Not Modified is NOT in the redirect frozenset → size check applies."""
        assert 304 not in _REDIRECT_STATUS_CODES
        oversized = _mock_response(
            304,
            {
                "Location": "https://other.com",
                "Content-Length": str(MAX_RESPONSE_SIZE + 1),
            },
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request(), allow_redirects=True)

    def test_oversized_intermediate_redirect_rejected(self, _bypass_ssrf):
        """301 + Location + allow_redirects=True + oversized → raises.

        All responses are size-checked uniformly, including redirects.
        """
        oversized = _mock_response(
            301,
            {
                "Location": "https://other.com",
                "Content-Length": str(MAX_RESPONSE_SIZE + 1),
            },
        )
        session = SafeSession()
        with patch.object(requests.Session, "send", return_value=oversized):
            with pytest.raises(ValueError, match="Response too large"):
                session.send(_prep_request(), allow_redirects=True)