Retry the token endpoint's POST on a read error

msal/application.py

@@ -267,6 +267,39 @@ def _msal_extension_check():
             )
 
 
+@functools.lru_cache(maxsize=1)
+def _retry_class(retry_base):
+    """Return a Retry subclass that retries errors on every method, POST included.
+
+    urllib3 retries a read error - a connection reset in the middle of a request, say -
+    only on the methods it considers idempotent, and POST is not one of them. The token
+    endpoint is the only thing MSAL POSTs, so the retry mounted below never applied to
+    it. Status-code retries stay on urllib3's default method set, so a 429 or a 5xx
+    carrying Retry-After still reaches MSAL's own throttling layer, rather than being
+    slept on and retried down here.
+    """
+    default_methods = getattr(
+        retry_base, "DEFAULT_ALLOWED_METHODS", None  # urllib3 1.26+
+        ) or getattr(retry_base, "DEFAULT_METHOD_WHITELIST", frozenset())
+
+    class RetryErrorsOnEveryMethod(retry_base):
+        def is_retry(self, method, status_code, has_retry_after=False):
+            if method is not None and method.upper() not in default_methods:
+                return False
+            return retry_base.is_retry(self, method, status_code, has_retry_after)
+
+    return RetryErrorsOnEveryMethod
+
+
+def _build_retry(requests):
+    """One retry for connection and read errors, on every method."""
+    retry_class = _retry_class(requests.adapters.Retry)
+    try:
+        return retry_class(total=1, allowed_methods=None)  # urllib3 1.26+
+    except TypeError:  # urllib3 < 1.26 spelled it method_whitelist
+        return retry_class(total=1, method_whitelist=None)
+
+
 class ClientApplication(object):
     """You do not usually directly use this class. Use its subclasses instead:
     :class:`PublicClientApplication` and :class:`ConfidentialClientApplication`.
@@ -715,7 +748,7 @@ def get_client_assertion():
 
             # Enable a minimal retry. Better than nothing.
             # https://github.com/psf/requests/blob/v2.25.1/requests/adapters.py#L94-L108
-            a = requests.adapters.HTTPAdapter(max_retries=1)
+            a = requests.adapters.HTTPAdapter(max_retries=_build_retry(requests))
             self.http_client.mount("http://", a)
             self.http_client.mount("https://", a)
         self.http_client = ThrottledHttpClient(

tests/test_http_retry.py

@@ -0,0 +1,32 @@
+import requests
+
+from msal.application import _build_retry
+
+from tests import unittest
+
+
+class TestRetryPolicy(unittest.TestCase):
+    """The retry mounted by ClientApplication has to cover the token endpoint's POST,
+    without taking over the 429/5xx handling that MSAL's ThrottledHttpClient does."""
+
+    def setUp(self):
+        self.retry = _build_retry(requests)
+
+    def test_one_retry(self):
+        self.assertEqual(1, self.retry.total)
+
+    def test_errors_are_retried_on_post(self):
+        # urllib3 gates read-error retries on whether it considers the method idempotent
+        self.assertTrue(self.retry._is_method_retryable("POST"))
+        self.assertTrue(self.retry._is_method_retryable("GET"))
+
+    def test_a_throttled_post_is_left_to_msal(self):
+        # A 429 or 5xx with Retry-After on the token endpoint reaches ThrottledHttpClient,
+        # instead of urllib3 sleeping for Retry-After seconds and retrying in place
+        self.assertFalse(self.retry.is_retry("POST", 429, has_retry_after=True))
+        self.assertFalse(self.retry.is_retry("POST", 503, has_retry_after=True))
+
+    def test_a_throttled_get_still_retries(self):
+        # Unchanged from the previous HTTPAdapter(max_retries=1): the discovery GETs
+        # are not throttled by MSAL itself
+        self.assertTrue(self.retry.is_retry("GET", 429, has_retry_after=True))