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