Implement claims challenge error for AzurePowerShellCredential

Co-authored-by: xiangyan99 <[email protected]>

Author: Copilot

sdk/identity/azure-identity/azure/identity/_credentials/azure_powershell.py

@@ -112,7 +112,7 @@ def close(self) -> None:
     def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -136,6 +136,12 @@ def get_token(
           receive an access token
         """
 
+        # Check if claims challenge is provided
+        if claims:
+            raise CredentialUnavailableError(
+                message="Fail to get token, please run Connect-AzAccount --ClaimsChallenge"
+            )
+
         options: TokenRequestOptions = {}
         if tenant_id:
             options["tenant_id"] = tenant_id
@@ -170,6 +176,12 @@ def _get_token_base(
         self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
     ) -> AccessTokenInfo:
 
+        # Check if claims challenge is provided
+        if options and options.get("claims"):
+            raise CredentialUnavailableError(
+                message="Fail to get token, please run Connect-AzAccount --ClaimsChallenge"
+            )
+
         tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)

sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_powershell.py

@@ -58,7 +58,7 @@ def __init__(
     async def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,  # pylint:disable=unused-argument
+        claims: Optional[str] = None,
         tenant_id: Optional[str] = None,
         **kwargs: Any,
     ) -> AccessToken:
@@ -80,6 +80,13 @@ async def get_token(
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked Azure PowerShell but didn't
           receive an access token
         """
+
+        # Check if claims challenge is provided
+        if claims:
+            raise CredentialUnavailableError(
+                message="Fail to get token, please run Connect-AzAccount --ClaimsChallenge"
+            )
+
         # only ProactorEventLoop supports subprocesses on Windows (and it isn't the default loop on Python < 3.8)
         if sys.platform.startswith("win") and not isinstance(asyncio.get_event_loop(), asyncio.ProactorEventLoop):
             return _SyncCredential().get_token(*scopes, tenant_id=tenant_id, **kwargs)
@@ -112,13 +119,27 @@ async def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptio
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked Azure PowerShell but didn't
           receive an access token
         """
+
+        # Check if claims challenge is provided
+        if options and options.get("claims"):
+            raise CredentialUnavailableError(
+                message="Fail to get token, please run Connect-AzAccount --ClaimsChallenge"
+            )
+
         if sys.platform.startswith("win") and not isinstance(asyncio.get_event_loop(), asyncio.ProactorEventLoop):
             return _SyncCredential().get_token_info(*scopes, options=options)
         return await self._get_token_base(*scopes, options=options)
 
     async def _get_token_base(
         self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
     ) -> AccessTokenInfo:
+
+        # Check if claims challenge is provided
+        if options and options.get("claims"):
+            raise CredentialUnavailableError(
+                message="Fail to get token, please run Connect-AzAccount --ClaimsChallenge"
+            )
+
         tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)

sdk/identity/azure-identity/tests/test_powershell_credential.py

@@ -380,3 +380,52 @@ def fake_Popen(command, **_):
                 kwargs = {"options": kwargs}
             token = getattr(credential, get_token_method)("scope", **kwargs)
             assert token.token == expected_token
+
+
+@pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
+def test_claims_challenge_error(get_token_method):
+    """The credential should raise CredentialUnavailableError when claims challenge is provided"""
+    
+    if get_token_method == "get_token":
+        # Test claims parameter in get_token method
+        with pytest.raises(CredentialUnavailableError) as exc_info:
+            getattr(AzurePowerShellCredential(), get_token_method)("scope", claims="some-claims")
+        assert "Fail to get token, please run Connect-AzAccount --ClaimsChallenge" in str(exc_info.value)
+    else:
+        # Test claims in options for get_token_info method
+        with pytest.raises(CredentialUnavailableError) as exc_info:
+            getattr(AzurePowerShellCredential(), get_token_method)("scope", options={"claims": "some-claims"})
+        assert "Fail to get token, please run Connect-AzAccount --ClaimsChallenge" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
+def test_empty_claims_no_error(get_token_method):
+    """The credential should not raise error for empty or None claims"""
+    
+    # Mock successful token response
+    expected_access_token = "access"
+    expected_expires_on = 1617923581
+    stdout = "azsdk%{}%{}".format(expected_access_token, expected_expires_on)
+    
+    Popen = get_mock_Popen(stdout=stdout)
+    with patch(POPEN, Popen):
+        if get_token_method == "get_token":
+            # Test None claims parameter in get_token method
+            token = getattr(AzurePowerShellCredential(), get_token_method)("scope", claims=None)
+            assert token.token == expected_access_token
+            
+            # Test empty string claims
+            token = getattr(AzurePowerShellCredential(), get_token_method)("scope", claims="")
+            assert token.token == expected_access_token
+        else:
+            # Test None claims in options for get_token_info method
+            token = getattr(AzurePowerShellCredential(), get_token_method)("scope", options={"claims": None})
+            assert token.token == expected_access_token
+            
+            # Test empty string claims in options
+            token = getattr(AzurePowerShellCredential(), get_token_method)("scope", options={"claims": ""})
+            assert token.token == expected_access_token
+            
+            # Test missing claims key in options
+            token = getattr(AzurePowerShellCredential(), get_token_method)("scope", options={})
+            assert token.token == expected_access_token

sdk/identity/azure-identity/tests/test_powershell_credential_async.py

@@ -389,3 +389,52 @@ async def fake_exec(*args, **_):
                 kwargs = {"options": kwargs}
             token = await getattr(credential, get_token_method)("scope", **kwargs)
             assert token.token == expected_token
+
+
+@pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
+async def test_claims_challenge_error(get_token_method):
+    """The credential should raise CredentialUnavailableError when claims challenge is provided"""
+    
+    if get_token_method == "get_token":
+        # Test claims parameter in get_token method
+        with pytest.raises(CredentialUnavailableError) as exc_info:
+            await getattr(AzurePowerShellCredential(), get_token_method)("scope", claims="some-claims")
+        assert "Fail to get token, please run Connect-AzAccount --ClaimsChallenge" in str(exc_info.value)
+    else:
+        # Test claims in options for get_token_info method
+        with pytest.raises(CredentialUnavailableError) as exc_info:
+            await getattr(AzurePowerShellCredential(), get_token_method)("scope", options={"claims": "some-claims"})
+        assert "Fail to get token, please run Connect-AzAccount --ClaimsChallenge" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
+async def test_empty_claims_no_error(get_token_method):
+    """The credential should not raise error for empty or None claims"""
+    
+    # Mock successful token response
+    expected_access_token = "access"
+    expected_expires_on = 1617923581
+    stdout = "azsdk%{}%{}".format(expected_access_token, expected_expires_on)
+    
+    exec_mock = get_mock_exec(stdout=stdout)
+    with patch(CREATE_SUBPROCESS_EXEC, exec_mock):
+        if get_token_method == "get_token":
+            # Test None claims parameter in get_token method
+            token = await getattr(AzurePowerShellCredential(), get_token_method)("scope", claims=None)
+            assert token.token == expected_access_token
+            
+            # Test empty string claims
+            token = await getattr(AzurePowerShellCredential(), get_token_method)("scope", claims="")
+            assert token.token == expected_access_token
+        else:
+            # Test None claims in options for get_token_info method
+            token = await getattr(AzurePowerShellCredential(), get_token_method)("scope", options={"claims": None})
+            assert token.token == expected_access_token
+            
+            # Test empty string claims in options
+            token = await getattr(AzurePowerShellCredential(), get_token_method)("scope", options={"claims": ""})
+            assert token.token == expected_access_token
+            
+            # Test missing claims key in options
+            token = await getattr(AzurePowerShellCredential(), get_token_method)("scope", options={})
+            assert token.token == expected_access_token