Merge pull request #810 from HardMax71/fix/strict-typing-retry-handler-pipeline

fix(http): pass mypy --strict in the retry handler and pipeline and fix the bugs it shows

Author: baywet

packages/http/httpx/kiota_http/middleware/middleware.py

@@ -19,43 +19,40 @@ class MiddlewarePipeline():
     it here https://buffered.dev/middleware-python-requests/
     """
 
-    def __init__(self, transport: httpx.AsyncBaseTransport):
+    def __init__(self, transport: httpx.AsyncBaseTransport) -> None:
         super().__init__()
-        self._current_middleware = None
-        self._first_middleware = None
+        self._current_middleware: Optional[BaseMiddleware] = None
+        self._first_middleware: Optional[BaseMiddleware] = None
         self._transport = transport
 
-    def add_middleware(self, middleware):
-        if self._middleware_present():
-            self._current_middleware.next = middleware
-            self._current_middleware = middleware
-        else:
+    def add_middleware(self, middleware: "BaseMiddleware") -> None:
+        if self._current_middleware is None:
             self._first_middleware = middleware
-            self._current_middleware = self._first_middleware
+        else:
+            self._current_middleware.next = middleware
+        self._current_middleware = middleware
 
-    async def send(self, request):
+    async def send(self, request: httpx.Request) -> httpx.Response:
 
-        if self._middleware_present():
+        if self._first_middleware is not None:
             return await self._first_middleware.send(request, self._transport)
-        # No middleware in pipeline, delete request optoions from header and
-        # send the request
-        del request.headers['request_options']
+        # No middleware in pipeline, remove the request options and send the request
+        request.extensions.pop(REQUEST_OPTIONS_KEY, None)
         return await self._transport.handle_async_request(request)
 
-    def _middleware_present(self):
-        return self._current_middleware
-
 
 class BaseMiddleware():
     """Base class for middleware. Handles moving a Request to the next middleware in the pipeline.
     If the current middleware is the last one in the pipeline, it makes a network request
     """
 
-    def __init__(self):
-        self.next = None
-        self.parent_span = None
+    def __init__(self) -> None:
+        self.next: Optional[BaseMiddleware] = None
+        self.parent_span: Optional[trace.Span] = None
 
-    async def send(self, request, transport):
+    async def send(
+        self, request: httpx.Request, transport: httpx.AsyncBaseTransport
+    ) -> httpx.Response:
         if self.next is None:
             # Remove request options if there's no other middleware in the chain.
             if hasattr(request, "extensions") and isinstance(request.extensions, dict):
@@ -65,19 +62,15 @@ async def send(self, request, transport):
             return response
         return await self.next.send(request, transport)
 
-    def _create_observability_span(self, request, span_name: str) -> trace.Span:
+    def _create_observability_span(self, request: httpx.Request, span_name: str) -> trace.Span:
         """Gets the parent_span from the request options and creates a new span.    
         If no parent_span is found in the request, uses the parent_span in the
         object. If parent_span is None, current context will be used."""
-        _span = None
         options = None
         if hasattr(request, "extensions") and isinstance(request.extensions, dict):
             options = request.extensions.get(REQUEST_OPTIONS_KEY)
         if options and (parent_span := options.get("parent_span", None)):
             self.parent_span = parent_span
-            _context = trace.set_span_in_context(parent_span)
-            _span = tracer.start_span(span_name, _context)
-        if _span is None:
-            _context = trace.set_span_in_context(self.parent_span)
-            _span = tracer.start_span(span_name, _context)
-        return _span
+        if self.parent_span is None:
+            return tracer.start_span(span_name)
+        return tracer.start_span(span_name, trace.set_span_in_context(self.parent_span))

packages/http/httpx/kiota_http/middleware/redirect_handler.py

@@ -84,8 +84,7 @@ async def send(
                     max_redirect -= 1
                     if not self.increment(response, max_redirect, history[:]):
                         if max_redirect < 0:
-                            response.history = history
-                            exc = RedirectError(f"Too many redirects. {response.history}")
+                            exc = RedirectError(f"Too many redirects. {history}")
                             _redirect_span.record_exception(exc)
                             raise exc
                         break
@@ -100,7 +99,6 @@ async def send(
                 break
             finally:
                 _redirect_span.end()
-        response.history = history
 
         return response
 

packages/http/httpx/kiota_http/middleware/retry_handler.py

@@ -3,6 +3,7 @@
 import random
 import re
 from email.utils import parsedate_to_datetime
+from typing import Optional
 
 from kiota_abstractions.request_option import RequestOption
 from opentelemetry.semconv.attributes.http_attributes import HTTP_RESPONSE_STATUS_CODE
@@ -42,16 +43,18 @@ class RetryHandler(BaseMiddleware):
         ['HEAD', 'GET', 'POST', 'PUT', 'PATCH', 'DELETE', 'OPTIONS', 'QUERY']
     )
 
-    def __init__(self, options: RetryHandlerOption = RetryHandlerOption()) -> None:
+    def __init__(self, options: Optional[RetryHandlerOption] = None) -> None:
         super().__init__()
         self.allowed_methods: frozenset[str] = self.DEFAULT_ALLOWED_METHODS
         self.backoff_factor: float = self.DEFAULT_BACKOFF_FACTOR
         self.backoff_max: int = self.MAXIMUM_BACKOFF
-        self.options = options
-        self.respect_retry_after_header: bool = self.options.DEFAULT_SHOULD_RETRY  # type:ignore
+        self.options = RetryHandlerOption() if options is None else options
+        self.respect_retry_after_header: bool = self.options.DEFAULT_SHOULD_RETRY
         self.retry_on_status_codes: set[int] = self.DEFAULT_RETRY_STATUS_CODES
 
-    async def send(self, request: httpx.Request, transport: httpx.AsyncBaseTransport):
+    async def send(
+        self, request: httpx.Request, transport: httpx.AsyncBaseTransport
+    ) -> httpx.Response:
         """
         Sends the http request object to the next middleware or retries the request if necessary.
         """
@@ -88,6 +91,7 @@ async def send(self, request: httpx.Request, transport: httpx.AsyncBaseTransport
                 request.headers.update({RETRY_ATTEMPT: f'{retry_count}'})
                 _retry_span.set_attribute(HTTP_RESPONSE_STATUS_CODE, response.status_code)
                 _retry_span.set_attribute('http.request.resend_count', retry_count)
+                _retry_span.end()
                 continue
             _retry_span.end()
             break
@@ -107,11 +111,15 @@ def _get_current_options(self, request: httpx.Request) -> RetryHandlerOption:
         """
         request_options = request.extensions.get(REQUEST_OPTIONS_KEY)
         if request_options:
-            current_options = request_options.get(RetryHandlerOption.get_key(), self.options)
+            current_options: RetryHandlerOption = request_options.get(
+                RetryHandlerOption.get_key(), self.options
+            )
             return current_options
         return self.options
 
-    def should_retry(self, request, options, response):
+    def should_retry(
+        self, request: httpx.Request, options: RetryHandlerOption, response: httpx.Response
+    ) -> bool:
         """
         Determines whether the request should be retried
         Checks if the request method is in allowed methods
@@ -121,10 +129,9 @@ def should_retry(self, request, options, response):
             return False
         if not self._is_request_payload_buffered(request):
             return False
-        val = options.max_retry and (response.status_code in self.retry_on_status_codes)
-        return val
+        return options.max_retry > 0 and response.status_code in self.retry_on_status_codes
 
-    def _is_method_retryable(self, request):
+    def _is_method_retryable(self, request: httpx.Request) -> bool:
         """
         Checks if a given request should be retried upon, depending on
         whether the HTTP method is in the set of allowed methods
@@ -133,7 +140,7 @@ def _is_method_retryable(self, request):
             return False
         return True
 
-    def _is_request_payload_buffered(self, request):
+    def _is_request_payload_buffered(self, request: httpx.Request) -> bool:
         """
         Checks if the request payload is buffered/rewindable.
         Payloads with forward only streams will return false and have the responses
@@ -145,15 +152,20 @@ def _is_request_payload_buffered(self, request):
             return False
         return True
 
-    def check_retry_valid(self, retry_count, options):
+    def check_retry_valid(self, retry_count: int, options: RetryHandlerOption) -> bool:
         """
         Check that the max retries limit has not been hit
         """
         if retry_count < options.max_retry:
             return True
         return False
 
-    def get_delay_time(self, retry_count, response=None, delay=RetryHandlerOption.DEFAULT_DELAY):
+    def get_delay_time(
+        self,
+        retry_count: int,
+        response: Optional[httpx.Response] = None,
+        delay: float = RetryHandlerOption.DEFAULT_DELAY
+    ) -> float:
         """Get the time in seconds to delay before the next attempt.
 
         A ``Retry-After`` response header that parses to more than zero seconds is
@@ -167,37 +179,41 @@ def get_delay_time(self, retry_count, response=None, delay=RetryHandlerOption.DE
             return retry_after
         return self._get_delay_time_exp_backoff(retry_count, delay)
 
-    def _get_delay_time_exp_backoff(self, retry_count, delay):
+    def _get_delay_time_exp_backoff(self, retry_count: int, delay: float) -> float:
         """
         Get time in seconds to delay between retry attempts based on an exponential
         backoff value.
         """
-        exp_backoff_value = self.backoff_factor * +(2**(retry_count - 1))
+        exp_backoff_value: float = self.backoff_factor * +(2**(retry_count - 1))
         backoff_value = exp_backoff_value + (random.randint(0, 1000) / 1000) + delay
 
         backoff = min(self.backoff_max, backoff_value)
         return backoff
 
-    def _get_retry_after(self, response):
+    def _get_retry_after(self, response: Optional[httpx.Response]) -> Optional[float]:
         """
         Check if retry-after is specified in the response header and get the value
         """
+        if response is None:
+            return None
         retry_after = response.headers.get("retry-after")
         if retry_after:
             return self._parse_retry_after(retry_after)
         return None
 
-    def _parse_retry_after(self, retry_after):
+    def _parse_retry_after(self, retry_after: str) -> Optional[float]:
         """
-        Helper to parse Retry-After and get value in seconds.
+        Helper to parse Retry-After and get value in seconds. Returns None when the
+        value is neither a number of seconds nor an HTTP date.
         """
+        if re.match(r"\d+,\d+$", retry_after):
+            retry_after = retry_after.split(",")[0]
         try:
-            retry_after = retry_after.split(",")[0] if re.match(
-                r"\d+,\d+$", retry_after
-            ) else retry_after
-            delay = int(retry_after)
+            return max(0, int(retry_after))
         except ValueError:
-            # Not an integer? Try HTTP date
+            pass
+        try:
             retry_date = parsedate_to_datetime(retry_after)
-            delay = (retry_date - datetime.datetime.now(retry_date.tzinfo)).total_seconds()
-        return max(0, delay)
+        except ValueError:
+            return None
+        return max(0.0, (retry_date - datetime.datetime.now(retry_date.tzinfo)).total_seconds())

packages/http/httpx/tests/middleware_tests/test_base_middleware.py

@@ -1,7 +1,10 @@
 """Test the BaseMiddleware class."""
+import httpx
+import pytest
 from opentelemetry import trace
 
-from kiota_http.middleware import BaseMiddleware
+from kiota_http.middleware import REQUEST_OPTIONS_KEY, BaseMiddleware, MiddlewarePipeline
+from kiota_http.middleware import middleware as middleware_module
 
 
 def test_next_is_none():
@@ -27,4 +30,24 @@ def test_span_created(request_info, otel_tracer):
         span = middleware._create_observability_span(request_info, "test_span_created")
         span.end()
         after_state = outside_span.is_recording()
-        assert(before_state == after_state)
+        assert(before_state == after_state)
+
+
+def test_span_without_parent_span_uses_the_current_span(span_exporter):
+    middleware = BaseMiddleware()
+    with middleware_module.tracer.start_as_current_span("outer"):
+        middleware._create_observability_span(httpx.Request("GET", "https://example.com"), "inner").end()
+    inner, outer = span_exporter.get_finished_spans()
+    assert inner.parent.span_id == outer.context.span_id
+
+
+@pytest.mark.asyncio
+async def test_pipeline_without_middleware_sends_the_request():
+    def request_handler(request: httpx.Request):
+        return httpx.Response(200)
+
+    request = httpx.Request("GET", "https://example.com", extensions={REQUEST_OPTIONS_KEY: {}})
+    async with httpx.MockTransport(request_handler) as transport:
+        response = await MiddlewarePipeline(transport).send(request)
+    assert response.status_code == 200
+    assert REQUEST_OPTIONS_KEY not in request.extensions

packages/http/httpx/tests/middleware_tests/test_redirect_handler.py

@@ -37,7 +37,6 @@ def request_handler(request):
     attempts = redirects + 1 if should_redirect else 1
     assert len(requests) == attempts
     assert response.status_code == (200 if should_redirect else 302)
-    assert len(response.history) == attempts - 1
     spans = span_exporter.get_finished_spans()
     assert [span.name for span in spans] == ["RedirectHandler_send"] + [
         f"RedirectHandler_send - redirect {index}" for index in range(attempts)
@@ -124,6 +123,27 @@ def test_increment_redirects():
     assert handler.increment(response, handler.options.max_redirect, history=history)
 
 
+class SingleRedirectHandler(RedirectHandler):
+
+    def increment(self, response, max_redirect, history):
+        return len(history) == 0
+
+
+@pytest.mark.asyncio
+async def test_increment_override_stops_redirects():
+    requests = []
+
+    def request_handler(request):
+        requests.append(request)
+        return httpx.Response(FOUND, headers={LOCATION_HEADER: f"/redirect/{len(requests)}"})
+
+    async with httpx.MockTransport(request_handler) as transport:
+        response = await SingleRedirectHandler().send(httpx.Request("GET", BASE_URL), transport)
+
+    assert response.status_code == FOUND
+    assert [request.url.path for request in requests] == ["/", "/redirect/1"]
+
+
 def test_same_origin(mock_redirect_handler):
     origin1 = httpx.URL(BASE_URL)
     origin2 = httpx.URL(f"{BASE_URL}:443")

packages/http/httpx/tests/middleware_tests/test_retry_handler.py

@@ -345,3 +345,34 @@ def request_handler(request: httpx.Request):
     assert responses[0].is_closed and responses[1].is_closed
     # the returned response is left open for the caller to read
     assert not responses[2].is_closed
+
+
+@pytest.mark.asyncio
+async def test_retry_spans_end(span_exporter):
+    def request_handler(request: httpx.Request):
+        if RETRY_ATTEMPT in request.headers:
+            return httpx.Response(200)
+        return httpx.Response(SERVICE_UNAVAILABLE)
+
+    async with httpx.MockTransport(request_handler) as transport:
+        await RetryHandler(RetryHandlerOption(delay=0)).send(httpx.Request("GET", BASE_URL), transport)
+    assert [span.name for span in span_exporter.get_finished_spans()] == [
+        "RetryHandler_send", "RetryHandler_send - attempt 0", "RetryHandler_send - attempt 1"
+    ]
+
+
+def test_default_options_are_not_shared():
+    first, second = RetryHandler(), RetryHandler()
+    first.options.max_retry = 7
+    assert second.options.max_retry == RetryHandlerOption.DEFAULT_MAX_RETRIES
+
+
+def test_get_delay_time_without_response():
+    delay = RetryHandler().get_delay_time(1, delay=0)
+    assert 0.5 <= delay <= 1.5
+
+
+def test_get_retry_after_invalid_value():
+    response = httpx.Response(SERVICE_UNAVAILABLE, headers={RETRY_AFTER: "soon"})
+    assert RetryHandler()._get_retry_after(response) is None
+