Add live test; assert scope

sdk/keyvault/azure-keyvault-keys/tests/test_challenge_auth.py

@@ -17,8 +17,8 @@
 from devtools_testutils import recorded_by_proxy
 
 import pytest
-from azure.core.credentials import AccessToken
-from azure.core.exceptions import ServiceRequestError
+from azure.core.credentials import AccessToken, TokenCredential
+from azure.core.exceptions import ClientAuthenticationError, ServiceRequestError
 from azure.core.pipeline import Pipeline
 from azure.core.pipeline.policies import SansIOHTTPPolicy
 from azure.core.rest import HttpRequest
@@ -69,6 +69,32 @@ def test_multitenant_authentication(self, client, is_hsm, **kwargs):
         else:
             os.environ.pop("AZURE_TENANT_ID")
 
+    @pytest.mark.skip("Manual test for specific, CAE-enabled environments.")
+    @pytest.mark.live_test_only
+    def test_cae_live(self, **kwargs):
+        class CredentialWrapper(TokenCredential):
+            def __init__(self, credential):
+                self._credential = credential
+                self._claims = None
+
+            def get_token(self, *scopes, **kwargs):
+                assert kwargs["enable_cae"] == True
+                if kwargs.get("claims"):
+                    # We should only receive claims once; subsequent challenges should be returned to the caller
+                    assert self._claims is None
+                    self._claims = kwargs["claims"]
+                return self._credential.get_token(*scopes, **kwargs)
+
+        credential = self.get_credential(KeyClient)
+        wrapped = CredentialWrapper(credential)
+        client = KeyClient(vault_url=os.environ["AZURE_KEYVAULT_URL"], credential=wrapped)
+        try:
+            client.create_rsa_key("key-name")  # Basic request meant to just trigger CAE challenges
+        # Test environment may continuously return claims challenges; a second consecutive challenge will raise
+        except ClientAuthenticationError as e:
+            assert "continuous access evaluation" in str(e).lower()
+        assert wrapped._claims is not None  # Ensure we passed a claim to a token request
+
 def empty_challenge_cache(fn):
     @functools.wraps(fn)
     def wrapper(**kwargs):
@@ -588,9 +614,10 @@ def send(request):
                 return kv_challenge
             raise ValueError("unexpected request")
 
-        def get_token(*_, **kwargs):
+        def get_token(*scopes, **kwargs):
             assert kwargs.get("enable_cae") == True
             assert kwargs.get("tenant_id") == tenant
+            assert scopes[0] == resource + "/.default"
             # Response to KV challenge
             if Requests.count == 1:
                 assert kwargs.get("claims") == None
@@ -669,9 +696,10 @@ def send(request):
                 return claims_challenge
             raise ValueError("unexpected request")
 
-        def get_token(*_, **kwargs):
+        def get_token(*scopes, **kwargs):
             assert kwargs.get("enable_cae") == True
             assert kwargs.get("tenant_id") == tenant
+            assert scopes[0] == resource + "/.default"
             # Response to KV challenge
             if Requests.count == 1:
                 assert kwargs.get("claims") == None