diff --git a/pyrit/auth/azure_auth.py b/pyrit/auth/azure_auth.py index b71616ec8e..9414ad687f 100644 --- a/pyrit/auth/azure_auth.py +++ b/pyrit/auth/azure_auth.py @@ -249,7 +249,8 @@ def __init__(self, token_scope: str, tenant_id: str = "") -> None: Args: token_scope (str): The token scope for authentication. - tenant_id (str, optional): The tenant ID. Defaults to "". + tenant_id (str, optional): The tenant ID to request the token for. If empty, the credential + chain uses its default tenant. Defaults to "". """ self._tenant_id = tenant_id self._token_scope = token_scope @@ -260,7 +261,10 @@ def _set_default_token(self) -> None: Set up default Azure credentials and retrieve access token. """ self.azure_creds = DefaultAzureCredential() - self.access_token = self.azure_creds.get_token(self._token_scope) + # DefaultAzureCredential rejects a tenant_id constructor argument, but get_token accepts one + # and forwards it to each credential in the chain. + token_kwargs = {"tenant_id": self._tenant_id} if self._tenant_id else {} + self.access_token = self.azure_creds.get_token(self._token_scope, **token_kwargs) self.token = self.access_token.token def refresh_token(self) -> str: diff --git a/tests/unit/auth/test_azure_auth.py b/tests/unit/auth/test_azure_auth.py index 0fe720af43..bf13b0bd46 100644 --- a/tests/unit/auth/test_azure_auth.py +++ b/tests/unit/auth/test_azure_auth.py @@ -51,6 +51,22 @@ def test_get_token_on_init() -> None: assert test_instance.token == mock_token +def test_get_token_forwards_tenant_id() -> None: + with patch( + "azure.identity.DefaultAzureCredential.get_token", return_value=AccessToken(mock_token, _TOKEN_TIME) + ) as mock_get_token: + AzureAuth(token_scope="https://mocked_endpoint.azure.com", tenant_id="my-tenant") + mock_get_token.assert_called_once_with("https://mocked_endpoint.azure.com", tenant_id="my-tenant") + + +def test_get_token_omits_tenant_id_when_not_set() -> None: + with patch( + "azure.identity.DefaultAzureCredential.get_token", return_value=AccessToken(mock_token, _TOKEN_TIME) + ) as mock_get_token: + AzureAuth(token_scope="https://mocked_endpoint.azure.com") + mock_get_token.assert_called_once_with("https://mocked_endpoint.azure.com") + + @pytest.mark.parametrize("elapsed_seconds", [0, 301]) def test_refresh_no_expiration(elapsed_seconds: int) -> None: with (