DI migrate tests (#36239)

xiangyan99 · web-flow · commit 3f37036f96c0 · 2024-07-22T15:24:24.000-07:00
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/dev_requirements.txt b/sdk/documentintelligence/azure-ai-documentintelligence/dev_requirements.txt
@@ -1,3 +1,4 @@
 -e ../../../tools/azure-sdk-tools
 ../../core/azure-core
-aiohttp
+aiohttp
+azure-identity
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/preparers.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/preparers.py
@@ -7,8 +7,7 @@
 
 import os
 import functools
-from devtools_testutils import PowerShellPreparer, AzureMgmtPreparer
-from azure.core.credentials import AzureKeyCredential
+from devtools_testutils import PowerShellPreparer, AzureMgmtPreparer, get_credential
 
 ENABLE_LOGGER = os.getenv("ENABLE_LOGGER", "False")
 REGION = os.getenv("DOCUMENTINTELLIGENCE_LOCATION", None)
@@ -18,7 +17,6 @@
     PowerShellPreparer,
     "documentintelligence",
     documentintelligence_endpoint="https://fakeendpoint.cognitiveservices.azure.com",
-    documentintelligence_api_key="fakeZmFrZV9hY29jdW50X2tleQ==",
     documentintelligence_storage_container_sas_url="https://blob_sas_url",
     documentintelligence_testing_data_container_sas_url="https://blob_sas_url",
     documentintelligence_training_data_classifier_sas_url="https://blob_sas_url",
@@ -36,16 +34,38 @@ def __init__(self, client_cls, client_kwargs={}, **kwargs):
     def create_resource(self, name, **kwargs):
         if self.is_live:
             documentintelligence_endpoint = os.environ["DOCUMENTINTELLIGENCE_ENDPOINT"]
-            documentintelligence_api_key = os.environ["DOCUMENTINTELLIGENCE_API_KEY"]
             polling_interval = 5
         else:
             documentintelligence_endpoint = "https://fakeendpoint.cognitiveservices.azure.com"
-            documentintelligence_api_key = "fakeZmFrZV9hY29jdW50X2tleQ=="
             polling_interval = 0
 
         client = self.client_cls(
             documentintelligence_endpoint,
-            AzureKeyCredential(documentintelligence_api_key),
+            get_credential(),
+            polling_interval=polling_interval,
+            logging_enable=True if ENABLE_LOGGER == "True" else False,
+            **self.client_kwargs
+        )
+        kwargs.update({"client": client})
+        return kwargs
+
+class GlobalClientPreparerAsync(AzureMgmtPreparer):
+    def __init__(self, client_cls, client_kwargs={}, **kwargs):
+        super(GlobalClientPreparerAsync, self).__init__(name_prefix="", random_name_length=42)
+        self.client_kwargs = client_kwargs
+        self.client_cls = client_cls
+
+    def create_resource(self, name, **kwargs):
+        if self.is_live:
+            documentintelligence_endpoint = os.environ["DOCUMENTINTELLIGENCE_ENDPOINT"]
+            polling_interval = 5
+        else:
+            documentintelligence_endpoint = "https://fakeendpoint.cognitiveservices.azure.com"
+            polling_interval = 0
+
+        client = self.client_cls(
+            documentintelligence_endpoint,
+            get_credential(is_async=True),
             polling_interval=polling_interval,
             logging_enable=True if ENABLE_LOGGER == "True" else False,
             **self.client_kwargs
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_custom_model.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_custom_model.py
@@ -6,8 +6,7 @@
 
 import pytest
 import uuid
-from devtools_testutils import recorded_by_proxy, set_bodiless_matcher
-from azure.core.credentials import AzureKeyCredential
+from devtools_testutils import recorded_by_proxy, set_bodiless_matcher, get_credential
 from azure.core.exceptions import ResourceNotFoundError
 from azure.ai.documentintelligence import DocumentIntelligenceClient, DocumentIntelligenceAdministrationClient
 from azure.ai.documentintelligence.models import (
@@ -24,9 +23,8 @@ class TestDACAnalyzeCustomModel(DocumentIntelligenceTest):
     @DocumentIntelligencePreparer()
     def test_analyze_document_none_model_id(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
         with pytest.raises(ValueError) as e:
             client.begin_analyze_document(model_id=None, analyze_request=b"xx")
@@ -35,9 +33,8 @@ def test_analyze_document_none_model_id(self, **kwargs):
     @DocumentIntelligencePreparer()
     def test_analyze_document_none_model_id_from_url(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
         with pytest.raises(ValueError) as e:
             client.begin_analyze_document(
@@ -49,9 +46,8 @@ def test_analyze_document_none_model_id_from_url(self, **kwargs):
     @recorded_by_proxy
     def test_analyze_document_empty_model_id(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
         with pytest.raises(ResourceNotFoundError) as e:
             client.begin_analyze_document(model_id="", analyze_request=b"xx")
@@ -62,9 +58,8 @@ def test_analyze_document_empty_model_id(self, **kwargs):
     @recorded_by_proxy
     def test_analyze_document_empty_model_id_from_url(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
         with pytest.raises(ResourceNotFoundError) as e:
             client.begin_analyze_document(
@@ -78,12 +73,11 @@ def test_analyze_document_empty_model_id_from_url(self, **kwargs):
     def test_custom_document_transform(self, documentintelligence_storage_container_sas_url, **kwargs):
         set_bodiless_matcher()
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         di_admin_client = DocumentIntelligenceAdministrationClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
         di_client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
 
         recorded_variables = kwargs.pop("variables", {})
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_custom_model_async.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_custom_model_async.py
@@ -7,8 +7,7 @@
 import pytest
 import uuid
 from devtools_testutils.aio import recorded_by_proxy_async
-from devtools_testutils import set_bodiless_matcher
-from azure.core.credentials import AzureKeyCredential
+from devtools_testutils import set_bodiless_matcher, get_credential
 from azure.core.exceptions import ResourceNotFoundError
 from azure.ai.documentintelligence.aio import DocumentIntelligenceClient, DocumentIntelligenceAdministrationClient
 from azure.ai.documentintelligence.models import (
@@ -26,9 +25,8 @@ class TestDACAnalyzeCustomModelAsync(AsyncDocumentIntelligenceTest):
     @recorded_by_proxy_async
     async def test_analyze_document_none_model_id(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
         with pytest.raises(ValueError) as e:
             async with client:
@@ -41,9 +39,8 @@ async def test_analyze_document_none_model_id(self, **kwargs):
     @recorded_by_proxy_async
     async def test_analyze_document_none_model_id_from_url(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
         with pytest.raises(ValueError) as e:
             async with client:
@@ -56,9 +53,8 @@ async def test_analyze_document_none_model_id_from_url(self, **kwargs):
     @recorded_by_proxy_async
     async def test_analyze_document_empty_model_id(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
         with pytest.raises(ResourceNotFoundError) as e:
             async with client:
@@ -72,9 +68,8 @@ async def test_analyze_document_empty_model_id(self, **kwargs):
     @recorded_by_proxy_async
     async def test_analyze_document_empty_model_id_from_url(self, **kwargs):
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
         with pytest.raises(ResourceNotFoundError) as e:
             async with client:
@@ -89,12 +84,11 @@ async def test_analyze_document_empty_model_id_from_url(self, **kwargs):
     async def test_custom_document_transform(self, documentintelligence_storage_container_sas_url, **kwargs):
         set_bodiless_matcher()
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         di_admin_client = DocumentIntelligenceAdministrationClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
         di_client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
 
         recorded_variables = kwargs.pop("variables", {})
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_layout.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_layout.py
@@ -6,8 +6,7 @@
 
 import pytest
 import functools
-from devtools_testutils import recorded_by_proxy
-from azure.core.credentials import AzureKeyCredential
+from devtools_testutils import recorded_by_proxy, get_credential
 from azure.ai.documentintelligence import DocumentIntelligenceClient
 from azure.ai.documentintelligence.models import (
     DocumentAnalysisFeature,
@@ -217,14 +216,14 @@ def test_layout_url_barcode(self, client):
     @skip_flaky_test
     @DocumentIntelligencePreparer()
     @recorded_by_proxy
-    def test_polling_interval(self, documentintelligence_endpoint, documentintelligence_api_key, **kwargs):
+    def test_polling_interval(self, documentintelligence_endpoint, **kwargs):
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
         assert client._config.polling_interval == 1
 
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key), polling_interval=7
+            documentintelligence_endpoint, get_credential(), polling_interval=7
         )
         assert client._config.polling_interval == 7
         poller = client.begin_analyze_document(
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_layout_async.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_analyze_layout_async.py
@@ -6,8 +6,8 @@
 
 import pytest
 import functools
+from devtools_testutils import get_credential
 from devtools_testutils.aio import recorded_by_proxy_async
-from azure.core.credentials import AzureKeyCredential
 from azure.ai.documentintelligence.aio import DocumentIntelligenceClient
 from azure.ai.documentintelligence.models import (
     DocumentAnalysisFeature,
@@ -16,7 +16,7 @@
 )
 from asynctestcase import AsyncDocumentIntelligenceTest
 from conftest import skip_flaky_test
-from preparers import DocumentIntelligencePreparer, GlobalClientPreparer as _GlobalClientPreparer
+from preparers import DocumentIntelligencePreparer, GlobalClientPreparerAsync as _GlobalClientPreparer
 
 
 DocumentIntelligenceClientPreparer = functools.partial(_GlobalClientPreparer, DocumentIntelligenceClient)
@@ -228,14 +228,14 @@ async def test_layout_url_barcodes(self, client):
     @skip_flaky_test
     @DocumentIntelligencePreparer()
     @recorded_by_proxy_async
-    async def test_polling_interval(self, documentintelligence_endpoint, documentintelligence_api_key, **kwargs):
+    async def test_polling_interval(self, documentintelligence_endpoint, **kwargs):
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
         assert client._config.polling_interval == 1
 
         client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key), polling_interval=7
+            documentintelligence_endpoint, get_credential(is_async=True), polling_interval=7
         )
         assert client._config.polling_interval == 7
         async with client:
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_classify_document.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_classify_document.py
@@ -6,8 +6,7 @@
 
 import functools
 import uuid
-from devtools_testutils import set_bodiless_matcher, recorded_by_proxy
-from azure.core.credentials import AzureKeyCredential
+from devtools_testutils import set_bodiless_matcher, recorded_by_proxy, get_credential
 from azure.ai.documentintelligence import DocumentIntelligenceAdministrationClient, DocumentIntelligenceClient
 from azure.ai.documentintelligence.models import (
     ClassifierDocumentTypeDetails,
@@ -32,12 +31,11 @@ class TestDACClassifyDocumentAsync(DocumentIntelligenceTest):
     def test_classify_document(self, documentintelligence_training_data_classifier_sas_url, **kwargs):
         set_bodiless_matcher()
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         di_client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
         di_admin_client = DocumentIntelligenceAdministrationClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential()
         )
 
         recorded_variables = kwargs.pop("variables", {})
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_classify_document_async.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dac_classify_document_async.py
@@ -6,9 +6,8 @@
 
 import functools
 import uuid
-from devtools_testutils import set_bodiless_matcher
+from devtools_testutils import set_bodiless_matcher, get_credential
 from devtools_testutils.aio import recorded_by_proxy_async
-from azure.core.credentials import AzureKeyCredential
 from azure.ai.documentintelligence.aio import DocumentIntelligenceAdministrationClient, DocumentIntelligenceClient
 from azure.ai.documentintelligence.models import (
     ClassifierDocumentTypeDetails,
@@ -18,7 +17,7 @@
 )
 from asynctestcase import AsyncDocumentIntelligenceTest
 from conftest import skip_flaky_test
-from preparers import DocumentIntelligencePreparer, GlobalClientPreparer as _GlobalClientPreparer
+from preparers import DocumentIntelligencePreparer, GlobalClientPreparerAsync as _GlobalClientPreparer
 
 
 DocumentModelAdministrationClientPreparer = functools.partial(
@@ -33,12 +32,11 @@ class TestDACClassifyDocumentAsync(AsyncDocumentIntelligenceTest):
     async def test_classify_document(self, documentintelligence_training_data_classifier_sas_url, **kwargs):
         set_bodiless_matcher()
         documentintelligence_endpoint = kwargs.pop("documentintelligence_endpoint")
-        documentintelligence_api_key = kwargs.pop("documentintelligence_api_key")
         di_client = DocumentIntelligenceClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
         di_admin_client = DocumentIntelligenceAdministrationClient(
-            documentintelligence_endpoint, AzureKeyCredential(documentintelligence_api_key)
+            documentintelligence_endpoint, get_credential(is_async=True)
         )
 
         recorded_variables = kwargs.pop("variables", {})
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dmac_classifiers_async.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dmac_classifiers_async.py
@@ -18,7 +18,7 @@
     BuildDocumentClassifierRequest,
 )
 from asynctestcase import AsyncDocumentIntelligenceTest
-from preparers import DocumentIntelligencePreparer, GlobalClientPreparer as _GlobalClientPreparer
+from preparers import DocumentIntelligencePreparer, GlobalClientPreparerAsync as _GlobalClientPreparer
 
 DocumentModelAdministrationClientPreparer = functools.partial(
     _GlobalClientPreparer, DocumentIntelligenceAdministrationClient
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dmac_compose_model_async.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dmac_compose_model_async.py
@@ -17,7 +17,7 @@
 )
 from asynctestcase import AsyncDocumentIntelligenceTest
 from conftest import skip_flaky_test
-from preparers import DocumentIntelligencePreparer, GlobalClientPreparer as _GlobalClientPreparer
+from preparers import DocumentIntelligencePreparer, GlobalClientPreparerAsync as _GlobalClientPreparer
 
 DocumentModelAdministrationClientPreparer = functools.partial(
     _GlobalClientPreparer, DocumentIntelligenceAdministrationClient
diff --git a/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dmac_copy_model_async.py b/sdk/documentintelligence/azure-ai-documentintelligence/tests/test_dmac_copy_model_async.py
@@ -18,7 +18,7 @@
 )
 from asynctestcase import AsyncDocumentIntelligenceTest
 from conftest import skip_flaky_test
-from preparers import DocumentIntelligencePreparer, GlobalClientPreparer as _GlobalClientPreparer
+from preparers import DocumentIntelligencePreparer, GlobalClientPreparerAsync as _GlobalClientPreparer
 
 DocumentModelAdministrationClientPreparer = functools.partial(
     _GlobalClientPreparer, DocumentIntelligenceAdministrationClient
diff --git a/sdk/documentintelligence/tests.yml b/sdk/documentintelligence/tests.yml
@@ -6,6 +6,7 @@ extends:
       BuildTargetingString: azure-ai-documentintelligence
       ServiceDirectory: documentintelligence
       TestTimeoutInMinutes: 200
+      UseFederatedAuth: true
       MatrixReplace:
         - TestSamples=.*/true
       Clouds: 'Prod'
@@ -15,12 +16,12 @@ extends:
       CloudConfig:
         Prod:
           SubscriptionConfiguration: $(sub-config-azure-cloud-test-resources-preview)
-          Location: 'eastus'
+          Location: eastus
+          ServiceConnection: azure-sdk-tests
+          SubscriptionConfigurationFilePaths: 
+            - eng/common/TestResources/sub-config/AzurePublicMsft.json
       EnvVars:
         AZURE_SUBSCRIPTION_ID: $(provisioner-subscription)
-        AZURE_TENANT_ID: $(aad-azure-sdk-test-tenant-id)
-        AZURE_CLIENT_SECRET: $(aad-azure-sdk-test-client-secret)
-        AZURE_CLIENT_ID: $(aad-azure-sdk-test-client-id)
         TEST_MODE: 'RunLiveNoRecord'
         AZURE_SKIP_LIVE_RECORDING: 'True'
         AZURE_TEST_RUN_LIVE: 'true'

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@`
`18`	`18`	`BuildDocumentClassifierRequest,`
`19`	`19`	`)`
`20`	`20`	`from asynctestcase import AsyncDocumentIntelligenceTest`
`21`		`-from preparers import DocumentIntelligencePreparer, GlobalClientPreparer as _GlobalClientPreparer`
	`21`	`+from preparers import DocumentIntelligencePreparer, GlobalClientPreparerAsync as _GlobalClientPreparer`
`22`	`22`
`23`	`23`	`DocumentModelAdministrationClientPreparer = functools.partial(`
`24`	`24`	`_GlobalClientPreparer, DocumentIntelligenceAdministrationClient`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@`
`17`	`17`	`)`
`18`	`18`	`from asynctestcase import AsyncDocumentIntelligenceTest`
`19`	`19`	`from conftest import skip_flaky_test`
`20`		`-from preparers import DocumentIntelligencePreparer, GlobalClientPreparer as _GlobalClientPreparer`
	`20`	`+from preparers import DocumentIntelligencePreparer, GlobalClientPreparerAsync as _GlobalClientPreparer`
`21`	`21`
`22`	`22`	`DocumentModelAdministrationClientPreparer = functools.partial(`
`23`	`23`	`_GlobalClientPreparer, DocumentIntelligenceAdministrationClient`