Handle connection errors in the https multi retriever

atlassian_jwt_auth/key.py

@@ -10,7 +10,7 @@
 import jwt
 import requests
 from cryptography.hazmat.primitives import serialization
-from requests.exceptions import RequestException
+from requests.exceptions import RequestException, ConnectionError
 
 from atlassian_jwt_auth.exceptions import (KeyIdentifierException,
                                            PublicKeyRetrieverException,
@@ -136,14 +136,24 @@ def _create_retrievers(self, key_repository_urls):
         return [HTTPSPublicKeyRetriever(url) for url
                 in key_repository_urls]
 
+    def handle_retrieval_exception(self, retriever, exception):
+        """ Handles working with exceptions encountered during key
+            retrieval.
+        """
+        if isinstance(exception, PublicKeyRetrieverException):
+            original_exception = getattr(
+                exception, 'original_exception', None)
+            if isinstance(original_exception, ConnectionError):
+                return
+            if exception.status_code is None or exception.status_code < 500:
+                raise
+
     def retrieve(self, key_identifier, **requests_kwargs):
         for retriever in self._retrievers:
             try:
                 return retriever.retrieve(key_identifier, **requests_kwargs)
             except (RequestException, PublicKeyRetrieverException) as e:
-                if isinstance(e, PublicKeyRetrieverException):
-                    if e.status_code is None or e.status_code < 500:
-                        raise
+                self.handle_retrieval_exception(retriever, e)
                 logger = logging.getLogger(__name__)
                 logger.warn('Unable to retrieve public key from store',
                             extra={'underlying_error': str(e),

atlassian_jwt_auth/tests/test_public_key_provider.py

@@ -190,6 +190,23 @@ def test_retrieve_with_500_error(self, mock_get_method):
             retriever.retrieve('example/eg'),
             self._public_key_pem)
 
+    @mock.patch.object(requests.Session, 'get')
+    def test_retrieve_with_connection_error(self, mock_get_method):
+        """ tests that the retrieve method works as expected
+            when the first key repository encounters a connection error.
+        """
+        retriever = HTTPSMultiRepositoryPublicKeyRetriever(self.keystore_urls)
+        _setup_mock_response_for_retriever(
+            mock_get_method, self._public_key_pem)
+        valid_response = mock_get_method.return_value
+        del mock_get_method.return_value
+        connection_exception = requests.exceptions.ConnectionError(
+            response=mock.Mock(status_code=None))
+        mock_get_method.side_effect = [connection_exception, valid_response]
+        self.assertEqual(
+            retriever.retrieve('example/eg'),
+            self._public_key_pem)
+
 
 def _setup_mock_response_for_retriever(
         mock_method, public_key_pem, headers=None):