[Identity] Fix live ACI test (#40434)

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/identity/azure-identity/tests/integration/azure-kubernetes-service/app.py

@@ -4,47 +4,65 @@
 # ------------------------------------
 import os
 import asyncio
+import argparse
 
 from azure.identity import ManagedIdentityCredential
 from azure.identity.aio import ManagedIdentityCredential as AsyncManagedIdentityCredential
 from azure.storage.blob import BlobServiceClient
 from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
 
 
-def run_sync():
-    credential = ManagedIdentityCredential()
+def run_sync(identity_type="system"):
+    """Run synchronous authentication using the specified identity type.
+
+    :param str identity_type: The type of managed identity to use ("system" or "user")
+    """
+    if identity_type == "user" and os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID"):
+        credential = ManagedIdentityCredential(client_id=os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID"))
+        storage_name = os.environ.get("IDENTITY_STORAGE_NAME_USER_ASSIGNED", os.environ.get("IDENTITY_STORAGE_NAME"))
+    else:
+        # Default to system-assigned identity
+        credential = ManagedIdentityCredential()
+        storage_name = os.environ.get("IDENTITY_STORAGE_NAME")
+
+    if not storage_name:
+        print("Storage account name not found in environment variables")
+        return False
 
     client = BlobServiceClient(
-        account_url=f"https://{os.environ['IDENTITY_STORAGE_NAME']}.blob.core.windows.net",
+        account_url=f"https://{storage_name}.blob.core.windows.net",
         credential=credential,
     )
 
     containers = client.list_containers()
     for container in containers:
         print(container["name"])
 
-    if os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID") and os.environ.get(
-        "IDENTITY_STORAGE_NAME_USER_ASSIGNED"
-    ):
-        credential = ManagedIdentityCredential(client_id=os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID"))
-
-        client = BlobServiceClient(
-            account_url=f"https://{os.environ['IDENTITY_STORAGE_NAME_USER_ASSIGNED']}.blob.core.windows.net",
-            credential=credential,
-        )
+    print(f"Successfully acquired token with ManagedIdentityCredential (identity_type={identity_type})")
+    return True
 
-        containers = client.list_containers()
-        for container in containers:
-            print(container["name"])
 
-    print(f"Successfully acquired token with ManagedIdentityCredential")
+async def run_async(identity_type="system"):
+    """Run asynchronous authentication using the specified identity type.
 
+    :param str identity_type: The type of managed identity to use ("system" or "user")
+    """
+    if identity_type == "user" and os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID"):
+        credential = AsyncManagedIdentityCredential(
+            client_id=os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID")
+        )
+        storage_name = os.environ.get("IDENTITY_STORAGE_NAME_USER_ASSIGNED", os.environ.get("IDENTITY_STORAGE_NAME"))
+    else:
+        # Default to system-assigned identity
+        credential = AsyncManagedIdentityCredential()
+        storage_name = os.environ.get("IDENTITY_STORAGE_NAME")
 
-async def run_async():
-    credential = AsyncManagedIdentityCredential()
+    if not storage_name:
+        print("Storage account name not found in environment variables")
+        return False
 
     client = AsyncBlobServiceClient(
-        account_url=f"https://{os.environ['IDENTITY_STORAGE_NAME']}.blob.core.windows.net",
+        account_url=f"https://{storage_name}.blob.core.windows.net",
         credential=credential,
     )
 
@@ -54,29 +72,44 @@ async def run_async():
     await client.close()
     await credential.close()
 
-    if os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID") and os.environ.get(
-        "IDENTITY_STORAGE_NAME_USER_ASSIGNED"
-    ):
-        credential = AsyncManagedIdentityCredential(
-            client_id=os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID")
-        )
-
-        client = AsyncBlobServiceClient(
-            account_url=f"https://{os.environ['IDENTITY_STORAGE_NAME_USER_ASSIGNED']}.blob.core.windows.net",
-            credential=credential,
-        )
-
-        async for container in client.list_containers():
-            print(container["name"])
-
-        await client.close()
-        await credential.close()
-
-    print("Successfully acquired token with async ManagedIdentityCredential")
+    print(f"Successfully acquired token with async ManagedIdentityCredential (identity_type={identity_type})")
+    return True
 
 
 if __name__ == "__main__":
-    run_sync()
-    asyncio.run(run_async())
-
-    print("Passed!")
+    parser = argparse.ArgumentParser(description="Test managed identity authentication in AKS")
+    parser.add_argument(
+        "--identity-type",
+        choices=["system", "user", "both"],
+        default="both",
+        help="Type of managed identity to use (system, user, or both)",
+    )
+    args = parser.parse_args()
+
+    success = True
+
+    if args.identity_type in ["system", "both"]:
+        print("Testing with system-assigned managed identity:")
+        if not run_sync("system"):
+            success = False
+
+    if args.identity_type in ["user", "both"] and os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID"):
+        print("\nTesting with user-assigned managed identity:")
+        if not run_sync("user"):
+            success = False
+
+    if args.identity_type in ["system", "both"]:
+        print("\nTesting with async system-assigned managed identity:")
+        if not asyncio.run(run_async("system")):
+            success = False
+
+    if args.identity_type in ["user", "both"] and os.environ.get("IDENTITY_USER_DEFINED_IDENTITY_CLIENT_ID"):
+        print("\nTesting with async user-assigned managed identity:")
+        if not asyncio.run(run_async("user")):
+            success = False
+
+    if success:
+        print("\nPassed!")
+    else:
+        print("\nFailed!")
+        exit(1)

sdk/identity/azure-identity/tests/integration/test_azure_container_instance.py

@@ -23,9 +23,9 @@ def test_azure_container_instance(self):
 
         # Using "script" as a workaround for "az container exec" requiring a tty.
         # https://github.com/Azure/azure-cli/issues/17530
-        command = (
-            f"{az_path} container exec -g {resource_group} -n {container_instance_name} --exec-command 'python /app.py'"
-        )
+        # Note: Currently only testing user-assigned identity as system-assigned identity doesn't seem to work if both a user-assigned
+        # identity and system-assigned identity are assigned to the container instance.
+        command = f"{az_path} container exec -g {resource_group} -n {container_instance_name} --exec-command 'python /app.py --identity-type user'"
         output = run_command(
             [
                 "script",