From 2761aa4d0026b269af22f800bfc32dcb14afaa1a Mon Sep 17 00:00:00 2001 From: Charles Lowell Date: Wed, 9 Jun 2021 16:55:16 -0700 Subject: [PATCH] Async/BearerTokenCredentialPolicy consistently calls on_exception --- .../core/pipeline/policies/_authentication.py | 9 ++++++-- .../policies/_authentication_async.py | 9 ++++++-- .../async_tests/test_authentication_async.py | 23 +++++++++++++++++-- .../azure-core/tests/test_authentication.py | 19 +++++++++++++++ 4 files changed, 54 insertions(+), 6 deletions(-) diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py b/sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py index 251dc8a610bf..228e3fd20f58 100644 --- a/sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py +++ b/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 diff --git a/sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py b/sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py index 479ef9057571..76564320b742 100644 --- a/sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py +++ b/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 diff --git a/sdk/core/azure-core/tests/async_tests/test_authentication_async.py b/sdk/core/azure-core/tests/async_tests/test_authentication_async.py index 047c0be0a6a0..7230018aa37f 100644 --- a/sdk/core/azure-core/tests/async_tests/test_authentication_async.py +++ b/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,6 +188,7 @@ 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]) @@ -195,6 +196,24 @@ class TestException(Exception): 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() diff --git a/sdk/core/azure-core/tests/test_authentication.py b/sdk/core/azure-core/tests/test_authentication.py index e11e146507d0..de029e8ea352 100644 --- a/sdk/core/azure-core/tests/test_authentication.py +++ b/sdk/core/azure-core/tests/test_authentication.py @@ -225,6 +225,7 @@ 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]) @@ -232,6 +233,24 @@ class TestException(Exception): 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():