Add Key Vault API versions for sovereign clouds in smoke tests

common/smoketest/key_vault_base.py

@@ -1,17 +1,22 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
 import os
-from azure.identity import DefaultAzureCredential, KnownAuthorities
+from azure.identity import AzureAuthorityHosts, DefaultAzureCredential
+
 
 class KeyVaultBase:
     credential_type = DefaultAzureCredential
     host_alias_map = {
-        'AzureChinaCloud': KnownAuthorities.AZURE_CHINA,
-        'AzureGermanCloud': KnownAuthorities.AZURE_GERMANY,
-        'AzureUSGovernment': KnownAuthorities.AZURE_GOVERNMENT,
-        'AzureCloud': KnownAuthorities.AZURE_PUBLIC_CLOUD,
+        "AzureChinaCloud": (AzureAuthorityHosts.AZURE_CHINA, "2016-10-01"),
+        "AzureGermanCloud": (AzureAuthorityHosts.AZURE_GERMANY, "2016-10-01"),
+        "AzureUSGovernment": (AzureAuthorityHosts.AZURE_GOVERNMENT, "2016-10-01"),
+        "AzureCloud": (AzureAuthorityHosts.AZURE_PUBLIC_CLOUD, "7.1"),
     }
 
-    # Instantiate a default credential based on the credential_type
-    def get_default_credential(self, authority_host_alias=None):
-        alias = authority_host_alias or os.environ.get("AZURE_CLOUD")
-        authority_host = self.host_alias_map.get(alias, KnownAuthorities.AZURE_PUBLIC_CLOUD)
-        return self.credential_type(authority=authority_host)
+    def get_client_args(self, authority_host_alias=None):
+        alias = authority_host_alias or os.environ.get("AZURE_CLOUD", "AzureCloud")
+        authority_host, api_version = self.host_alias_map[alias]
+        credential = self.credential_type(authority=authority_host)
+        return {"api_version": api_version, "credential": credential, "vault_url": os.environ["AZURE_PROJECT_URL"]}

common/smoketest/key_vault_base_async.py

@@ -1,5 +1,10 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
 from key_vault_base import KeyVaultBase
 from azure.identity.aio import DefaultAzureCredential
 
+
 class KeyVaultBaseAsync(KeyVaultBase):
     credential_type = DefaultAzureCredential

common/smoketest/key_vault_certificates.py

@@ -2,19 +2,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import os
 import uuid
 from azure.keyvault.certificates import CertificateClient, CertificatePolicy
 from key_vault_base import KeyVaultBase
 
+
 class KeyVaultCertificates(KeyVaultBase):
     def __init__(self):
-
-        credential = self.get_default_credential()
-        self.certificate_client = CertificateClient(
-            vault_url=os.environ["AZURE_PROJECT_URL"], credential=credential
-        )
-
+        args = self.get_client_args()
+        self.certificate_client = CertificateClient(**args)
         self.certificate_name = "cert-name-" + uuid.uuid1().hex
 
     def create_certificate(self):

common/smoketest/key_vault_certificates_async.py

@@ -2,19 +2,16 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import os
 import uuid
 from azure.keyvault.certificates import CertificatePolicy
 from azure.keyvault.certificates.aio import CertificateClient
 from key_vault_base_async import KeyVaultBaseAsync
 
+
 class KeyVaultCertificates(KeyVaultBaseAsync):
     def __init__(self):
-        credential = self.get_default_credential()
-        self.certificate_client = CertificateClient(
-            vault_url=os.environ["AZURE_PROJECT_URL"], credential=credential
-        )
-
+        args = self.get_client_args()
+        self.certificate_client = CertificateClient(**args)
         self.certificate_name = "cert-name-" + uuid.uuid1().hex
 
     async def create_certificate(self):

common/smoketest/key_vault_keys.py

@@ -2,19 +2,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import os
 import uuid
 from azure.keyvault.keys import KeyClient
 from key_vault_base import KeyVaultBase
 
 
 class KeyVaultKeys(KeyVaultBase):
     def __init__(self):
-        credential = self.get_default_credential()
-        self.key_client = KeyClient(
-            vault_url=os.environ["AZURE_PROJECT_URL"], credential=credential
-        )
-
+        args = self.get_client_args()
+        self.key_client = KeyClient(**args)
         self.key_name = "key-name-" + uuid.uuid1().hex
 
     def create_rsa_key(self):

common/smoketest/key_vault_keys_async.py

@@ -2,20 +2,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import os
 import uuid
 from azure.keyvault.keys.aio import KeyClient
 from key_vault_base_async import KeyVaultBaseAsync
 
 
 class KeyVaultKeys(KeyVaultBaseAsync):
     def __init__(self):
-
-        credential = self.get_default_credential()
-        self.key_client = KeyClient(
-            vault_url=os.environ["AZURE_PROJECT_URL"], credential=credential
-        )
-
+        args = self.get_client_args()
+        self.key_client = KeyClient(**args)
         self.key_name = "key-name-" + uuid.uuid1().hex
 
     async def create_rsa_key(self):

common/smoketest/key_vault_secrets.py

@@ -2,18 +2,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import os
 import uuid
 from azure.keyvault.secrets import SecretClient
 from key_vault_base import KeyVaultBase
 
+
 class KeyVaultSecrets(KeyVaultBase):
     def __init__(self):
-        credential = self.get_default_credential()
-        self.secret_client = SecretClient(
-            vault_url=os.environ["AZURE_PROJECT_URL"], credential=credential
-        )
-
+        args = self.get_client_args()
+        self.secret_client = SecretClient(**args)
         self.secret_name = "secret-name-" + uuid.uuid1().hex
         self.secret_Value = "secret-value"
 

common/smoketest/key_vault_secrets_async.py

@@ -2,17 +2,15 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import os
 import uuid
 from azure.keyvault.secrets.aio import SecretClient
 from key_vault_base_async import KeyVaultBaseAsync
 
+
 class KeyVaultSecrets(KeyVaultBaseAsync):
     def __init__(self):
-        credential = self.get_default_credential()
-        self.secret_client = SecretClient(
-            vault_url=os.environ["AZURE_PROJECT_URL"], credential=credential
-        )
+        args = self.get_client_args()
+        self.secret_client = SecretClient(**args)
         self.secret_name = "secret-name-" + uuid.uuid1().hex
         self.secret_value = "secret-value"