"""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)