Async/BearerTokenCredentialPolicy consistently calls on_exception

sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py

@@ -128,8 +128,13 @@ def send(self, request):
                 if "WWW-Authenticate" in response.http_response.headers:
                     request_authorized = self.on_challenge(request, response)
                     if request_authorized:
-                        response = self.next.send(request)
-                        self.on_response(request, response)
+                        try:
+                            response = self.next.send(request)
+                            self.on_response(request, response)
+                        except Exception:  # pylint:disable=broad-except
+                            handled = self.on_exception(request)
+                            if not handled:
+                                raise
 
         return response
 

sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py

@@ -84,8 +84,13 @@ async def send(self, request: "PipelineRequest") -> "PipelineResponse":
                 if "WWW-Authenticate" in response.http_response.headers:
                     request_authorized = await self.on_challenge(request, response)
                     if request_authorized:
-                        response = await self.next.send(request)
-                        await await_result(self.on_response, request, response)
+                        try:
+                            response = await self.next.send(request)
+                            await await_result(self.on_response, request, response)
+                        except Exception:  # pylint:disable=broad-except
+                            handled = await await_result(self.on_exception, request)
+                            if not handled:
+                                raise
 
         return response
 

sdk/core/azure-core/tests/async_tests/test_authentication_async.py

@@ -164,14 +164,14 @@ async def test_bearer_policy_calls_sansio_methods():
 
     class TestPolicy(AsyncBearerTokenCredentialPolicy):
         def __init__(self, *args, **kwargs):
-            super(TestPolicy, self).__init__(*args, **kwargs)
+            super().__init__(*args, **kwargs)
             self.on_exception = Mock(return_value=False)
             self.on_request = Mock()
             self.on_response = Mock()
 
         async def send(self, request):
             self.request = request
-            self.response = await super(TestPolicy, self).send(request)
+            self.response = await super().send(request)
             return self.response
 
     credential = Mock(get_token=Mock(return_value=get_completed_future(AccessToken("***", int(time.time()) + 3600))))
@@ -188,13 +188,32 @@ async def send(self, request):
     class TestException(Exception):
         pass
 
+    # during the first send...
     transport = Mock(send=Mock(side_effect=TestException))
     policy = TestPolicy(credential, "scope")
     pipeline = AsyncPipeline(transport=transport, policies=[policy])
     with pytest.raises(TestException):
         await pipeline.run(HttpRequest("GET", "https://localhost"))
     policy.on_exception.assert_called_once_with(policy.request)
 
+    # ...or the second
+    async def fake_send(*args, **kwargs):
+        if fake_send.calls == 0:
+            fake_send.calls = 1
+            return Mock(status_code=401, headers={"WWW-Authenticate": 'Basic realm="localhost"'})
+        raise TestException()
+    fake_send.calls = 0
+
+    policy = TestPolicy(credential, "scope")
+    policy.on_challenge = Mock(return_value=get_completed_future(True))
+    transport = Mock(send=Mock(wraps=fake_send))
+    pipeline = AsyncPipeline(transport=transport, policies=[policy])
+    with pytest.raises(TestException):
+        await pipeline.run(HttpRequest("GET", "https://localhost"))
+    assert transport.send.call_count == 2
+    policy.on_challenge.assert_called_once()
+    policy.on_exception.assert_called_once_with(policy.request)
+
 
 def get_completed_future(result=None):
     fut = asyncio.Future()

sdk/core/azure-core/tests/test_authentication.py

@@ -225,13 +225,32 @@ def send(self, request):
     class TestException(Exception):
         pass
 
+    # during the first send...
     transport = Mock(send=Mock(side_effect=TestException))
     policy = TestPolicy(credential, "scope")
     pipeline = Pipeline(transport=transport, policies=[policy])
     with pytest.raises(TestException):
         pipeline.run(HttpRequest("GET", "https://localhost"))
     policy.on_exception.assert_called_once_with(policy.request)
 
+    # ...or the second
+    def raise_the_second_time(*args, **kwargs):
+        if raise_the_second_time.calls == 0:
+            raise_the_second_time.calls = 1
+            return Mock(status_code=401, headers={"WWW-Authenticate": 'Basic realm="localhost"'})
+        raise TestException()
+    raise_the_second_time.calls = 0
+
+    policy = TestPolicy(credential, "scope")
+    policy.on_challenge = Mock(return_value=True)
+    transport = Mock(send=Mock(wraps=raise_the_second_time))
+    pipeline = Pipeline(transport=transport, policies=[policy])
+    with pytest.raises(TestException):
+        pipeline.run(HttpRequest("GET", "https://localhost"))
+    assert transport.send.call_count == 2
+    policy.on_challenge.assert_called_once()
+    policy.on_exception.assert_called_once_with(policy.request)
+
 
 @pytest.mark.skipif(azure.core.__version__ >= "2", reason="this test applies only to azure-core 1.x")
 def test_key_vault_regression():