diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py b/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py index a94e001852ec..e443fc8fea53 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py @@ -36,6 +36,10 @@ class AzureCliCredential(object): This requires previously logging in to Azure via "az login", and will use the CLI's currently logged in identity. """ + def __init__(self, tenant_id: str = ""): + object.__init__(self) + + self.tenant_id = tenant_id def __enter__(self): return self @@ -67,7 +71,8 @@ def get_token(self, *scopes, **kwargs): # pylint: disable=no-self-use resource = _scopes_to_resource(*scopes) command = COMMAND_LINE.format(resource) - tenant = resolve_tenant("", **kwargs) + tenant = resolve_tenant(default_tenant= self.tenant_id, **kwargs) + if tenant: command += " --tenant " + tenant output = _run_command(command) diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py index 869cf5de69ae..33ee4f5cac96 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py @@ -32,6 +32,10 @@ class AzureCliCredential(AsyncContextManager): This requires previously logging in to Azure via "az login", and will use the CLI's currently logged in identity. """ + def __init__(self, tenant_id: str = ""): + AsyncContextManager.__init__(self) + + self.tenant_id = tenant_id @log_get_token_async async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken": @@ -55,7 +59,8 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken": resource = _scopes_to_resource(*scopes) command = COMMAND_LINE.format(resource) - tenant = resolve_tenant("", **kwargs) + tenant = resolve_tenant(default_tenant= self.tenant_id, **kwargs) + if tenant: command += " --tenant " + tenant output = await _run_command(command) diff --git a/sdk/identity/azure-identity/tests/test_cli_credential.py b/sdk/identity/azure-identity/tests/test_cli_credential.py index 6eb71b97e722..1ed04ae3e0b8 100644 --- a/sdk/identity/azure-identity/tests/test_cli_credential.py +++ b/sdk/identity/azure-identity/tests/test_cli_credential.py @@ -152,6 +152,36 @@ def test_timeout(): AzureCliCredential().get_token("scope") +def test_multitenant_authentication_class(): + default_tenant = "first-tenant" + first_token = "***" + second_tenant = "second-tenant" + second_token = first_token * 2 + + def fake_check_output(command_line, **_): + match = re.search("--tenant (.*)", command_line[-1]) + tenant = match.groups()[0] if match else default_tenant + assert tenant in (default_tenant, second_tenant), 'unexpected tenant "{}"'.format(tenant) + return json.dumps( + { + "expiresOn": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"), + "accessToken": first_token if tenant == default_tenant else second_token, + "subscription": "some-guid", + "tenant": tenant, + "tokenType": "Bearer", + } + ) + + with mock.patch(CHECK_OUTPUT, fake_check_output): + token = AzureCliCredential().get_token("scope") + assert token.token == first_token + + token = AzureCliCredential(tenant_id= default_tenant).get_token("scope") + assert token.token == first_token + + token = AzureCliCredential(tenant_id= second_tenant).get_token("scope") + assert token.token == second_token + def test_multitenant_authentication(): default_tenant = "first-tenant" first_token = "***"