Fix inability to call response.raise_for_status in AsyncTenacityTransport (#2668)

Author: dmontagu

pydantic_ai_slim/pydantic_ai/retries.py

@@ -28,7 +28,7 @@
 
 from datetime import datetime, timezone
 from email.utils import parsedate_to_datetime
-from typing import Callable, cast
+from typing import Any, Callable, cast
 
 from httpx import HTTPStatusError
 from tenacity import RetryCallState, wait_exponential
@@ -78,7 +78,7 @@ def __init__(
         self,
         controller: Retrying,
         wrapped: BaseTransport | None = None,
-        validate_response: Callable[[Response], None] | None = None,
+        validate_response: Callable[[Response], Any] | None = None,
     ):
         self.controller = controller
         self.wrapped = wrapped or HTTPTransport()
@@ -100,6 +100,10 @@ def handle_request(self, request: Request) -> Response:
         for attempt in self.controller:
             with attempt:
                 response = self.wrapped.handle_request(request)
+
+                # this is normally set by httpx _after_ calling this function, but we want the request in the validator:
+                response.request = request
+
                 if self.validate_response:
                     self.validate_response(response)
                 return response
@@ -149,7 +153,7 @@ def __init__(
         self,
         controller: AsyncRetrying,
         wrapped: AsyncBaseTransport | None = None,
-        validate_response: Callable[[Response], None] | None = None,
+        validate_response: Callable[[Response], Any] | None = None,
     ):
         self.controller = controller
         self.wrapped = wrapped or AsyncHTTPTransport()
@@ -171,6 +175,10 @@ async def handle_async_request(self, request: Request) -> Response:
         async for attempt in self.controller:
             with attempt:
                 response = await self.wrapped.handle_async_request(request)
+
+                # this is normally set by httpx _after_ calling this function, but we want the request in the validator:
+                response.request = request
+
                 if self.validate_response:
                     self.validate_response(response)
                 return response

tests/test_tenacity.py

@@ -134,6 +134,54 @@ def validate_response(response: httpx.Response):
         assert result is mock_response_success
         assert mock_transport.handle_request.call_count == 2
 
+    def test_raise_for_status_in_validate_response(self):
+        """Test that response.raise_for_status() works in validate_response callback."""
+        mock_transport = Mock(spec=httpx.BaseTransport)
+        mock_response_fail = Mock(spec=httpx.Response)
+        mock_response_fail.status_code = 429
+        mock_response_fail.is_success = False
+        mock_response_fail.is_error = True
+        mock_response_fail.request = None  # Initially None, will be set by transport
+
+        # Mock raise_for_status to check if request is set
+        def mock_raise_for_status():
+            if mock_response_fail.request is None:
+                raise RuntimeError(  # pragma: no cover
+                    'Cannot call `raise_for_status` as the request instance has not been set on this response.'
+                )
+            raise httpx.HTTPStatusError(
+                'Too Many Requests', request=mock_response_fail.request, response=mock_response_fail
+            )
+
+        mock_response_fail.raise_for_status = mock_raise_for_status
+
+        mock_response_success = Mock(spec=httpx.Response)
+        mock_response_success.status_code = 200
+        mock_response_success.is_success = True
+        mock_response_success.is_error = False
+        mock_response_success.raise_for_status = Mock()  # Should not raise
+
+        mock_transport.handle_request.side_effect = [mock_response_fail, mock_response_success]
+
+        controller = Retrying(
+            retry=retry_if_exception_type(httpx.HTTPStatusError),
+            stop=stop_after_attempt(3),
+            wait=wait_fixed(0.001),
+            reraise=True,
+        )
+        transport = TenacityTransport(
+            controller, mock_transport, validate_response=lambda response: response.raise_for_status()
+        )
+
+        request = httpx.Request('GET', 'https://example.com')
+        result = transport.handle_request(request)
+
+        assert result is mock_response_success
+        assert mock_transport.handle_request.call_count == 2
+        # Verify that the request was set on the failed response before raise_for_status was called
+        assert mock_response_fail.request is request
+        mock_response_success.raise_for_status.assert_called_once()
+
 
 class TestAsyncTenacityTransport:
     """Tests for the asynchronous AsyncTenacityTransport."""
@@ -243,6 +291,54 @@ def validate_response(response: httpx.Response):
         assert result is mock_response_success
         assert mock_transport.handle_async_request.call_count == 2
 
+    async def test_raise_for_status_in_validate_response(self):
+        """Test that response.raise_for_status() works in validate_response callback."""
+        mock_transport = AsyncMock(spec=httpx.AsyncBaseTransport)
+        mock_response_fail = Mock(spec=httpx.Response)
+        mock_response_fail.status_code = 429
+        mock_response_fail.is_success = False
+        mock_response_fail.is_error = True
+        mock_response_fail.request = None  # Initially None, will be set by transport
+
+        # Mock raise_for_status to check if request is set
+        def mock_raise_for_status():
+            if mock_response_fail.request is None:
+                raise RuntimeError(  # pragma: no cover
+                    'Cannot call `raise_for_status` as the request instance has not been set on this response.'
+                )
+            raise httpx.HTTPStatusError(
+                'Too Many Requests', request=mock_response_fail.request, response=mock_response_fail
+            )
+
+        mock_response_fail.raise_for_status = mock_raise_for_status
+
+        mock_response_success = Mock(spec=httpx.Response)
+        mock_response_success.status_code = 200
+        mock_response_success.is_success = True
+        mock_response_success.is_error = False
+        mock_response_success.raise_for_status = Mock()  # Should not raise
+
+        mock_transport.handle_async_request.side_effect = [mock_response_fail, mock_response_success]
+
+        controller = AsyncRetrying(
+            retry=retry_if_exception_type(httpx.HTTPStatusError),
+            stop=stop_after_attempt(3),
+            wait=wait_fixed(0.001),
+            reraise=True,
+        )
+        transport = AsyncTenacityTransport(
+            controller, mock_transport, validate_response=lambda response: response.raise_for_status()
+        )
+
+        request = httpx.Request('GET', 'https://example.com')
+        result = await transport.handle_async_request(request)
+
+        assert result is mock_response_success
+        assert mock_transport.handle_async_request.call_count == 2
+        # Verify that the request was set on the failed response before raise_for_status was called
+        assert mock_response_fail.request is request
+        mock_response_success.raise_for_status.assert_called_once()
+
 
 class TestWaitRetryAfter:
     """Tests for the wait_retry_after wait strategy."""