[ASYNC] apply #2281 to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_wif_util.py

@@ -202,7 +202,10 @@ async def create_azure_attestation(
     issuer, subject = extract_iss_and_sub_without_signature_verification(jwt_str)
     if not issuer or not subject:
         return None
-    if not issuer.startswith("https://sts.windows.net/"):
+    if not (
+        issuer.startswith("https://sts.windows.net/")
+        or issuer.startswith("https://login.microsoftonline.com/")
+    ):
         # This might happen if we're running on a different platform that responds to the same metadata request signature as Azure.
         logger.debug("Unexpected Azure token issuer '%s'", issuer)
         return None

test/unit/aio/test_auth_workload_identity_async.py

@@ -261,14 +261,34 @@ async def test_explicit_azure_metadata_server_error_raises_auth_error(exception)
 
 
 async def test_explicit_azure_wrong_issuer_raises_error(fake_azure_metadata_service):
-    fake_azure_metadata_service.iss = "not-azure"
+    fake_azure_metadata_service.iss = "https://notazure.com"
 
     auth_class = AuthByWorkloadIdentity(provider=AttestationProvider.AZURE)
     with pytest.raises(ProgrammingError) as excinfo:
         await auth_class.prepare()
     assert "No workload identity credential was found for 'AZURE'" in str(excinfo.value)
 
 
+@pytest.mark.parametrize(
+    "issuer",
+    [
+        "https://sts.windows.net/067802cd-8f92-4c7c-bceb-ea8f15d31cc5",
+        "https://login.microsoftonline.com/067802cd-8f92-4c7c-bceb-ea8f15d31cc5",
+        "https://login.microsoftonline.com/067802cd-8f92-4c7c-bceb-ea8f15d31cc5/v2.0",
+    ],
+    ids=["v1", "v2_without_suffix", "v2_with_suffix"],
+)
+async def test_explicit_azure_v1_and_v2_issuers_accepted(
+    fake_azure_metadata_service, issuer
+):
+    fake_azure_metadata_service.iss = issuer
+
+    auth_class = AuthByWorkloadIdentity(provider=AttestationProvider.AZURE)
+    await auth_class.prepare()
+
+    assert issuer == json.loads(auth_class.assertion_content)["iss"]
+
+
 async def test_explicit_azure_plumbs_token_to_api(fake_azure_metadata_service):
     auth_class = AuthByWorkloadIdentity(provider=AttestationProvider.AZURE)
     await auth_class.prepare()