From 4c236dde1b747245f0bbbd6135ef6246579eda99 Mon Sep 17 00:00:00 2001 From: Ming Xu Date: Sat, 28 Mar 2026 21:06:15 +1100 Subject: [PATCH] Implement subscription filter on az login --- src/azure-cli-core/azure/cli/core/_profile.py | 42 ++- .../azure/cli/core/tests/test_profile.py | 285 +++++++++++++++++- .../cli/command_modules/profile/__init__.py | 12 + .../cli/command_modules/profile/_help.py | 6 + .../cli/command_modules/profile/custom.py | 15 +- .../tests/latest/test_profile_custom.py | 72 +++++ 6 files changed, 422 insertions(+), 10 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 9b192b84199..b677217d5ac 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/src/azure-cli-core/azure/cli/core/_profile.py @@ -152,7 +152,9 @@ def login(self, allow_no_subscriptions=False, use_cert_sn_issuer=None, show_progress=False, - claims_challenge=None): + claims_challenge=None, + skip_subscription_discovery=False, + default_subscription=None): """ For service principal, `password` is a dict returned by ServicePrincipalAuth.build_credential """ @@ -198,15 +200,25 @@ def login(self, else: credential = identity.get_service_principal_credential(username) - if tenant: + is_bare_mode = skip_subscription_discovery and not default_subscription + + if skip_subscription_discovery and default_subscription: + # Fast path: fetch only the specified subscription (1 API call) + subscriptions = subscription_finder.find_specific_subscriptions( + tenant, credential, [default_subscription]) + elif is_bare_mode: + # Bare mode: no ARM subscription calls. Tenant-level account will be creatd below + subscriptions = [] + subscription_finder.tenants.append(tenant) + elif tenant: subscriptions = subscription_finder.find_using_specific_tenant(tenant, credential) else: subscriptions = subscription_finder.find_using_common_tenant(username, credential) - if not subscriptions and not allow_no_subscriptions: + if not subscriptions and not allow_no_subscriptions and not is_bare_mode: raise CLIError("No subscriptions found for {}.".format(username)) - if allow_no_subscriptions: + if allow_no_subscriptions or is_bare_mode: t_list = [s.tenant_id for s in subscriptions] bare_tenants = [t for t in subscription_finder.tenants if t not in t_list] tenant_accounts = self._build_tenant_level_accounts(bare_tenants) @@ -218,6 +230,27 @@ def login(self, is_service_principal, bool(use_cert_sn_issuer)) self._set_subscriptions(consolidated) + + # Validate default_subscription exists before calling set_active_subscription, + # so the caller can handle "not found" differently based on context. + if default_subscription: + match = next((s for s in consolidated + if s[_SUBSCRIPTION_ID].lower() == default_subscription.lower() or + s.get(_SUBSCRIPTION_NAME, '').lower() == default_subscription.lower()), None) + if match: + self.set_active_subscription(match[_SUBSCRIPTION_ID]) + # Refresh consolidated from storage so the returned list reflects the new default + consolidated = self.load_cached_subscriptions() + elif skip_subscription_discovery: + # --skip-subscription-discovery + --subscription S, but S is inaccessible. + # without --allow-no-subscriptions → already errored above + # with --allow-no-subscriptions → tenant-level account only (we reach here) + logger.warning("Subscription '%s' not found. Profile has tenant-level account only.", + default_subscription) + else: + raise CLIError("Subscription '{}' not found. Check the ID or name and try again." + .format(default_subscription)) + return deepcopy(consolidated) def login_with_managed_identity(self, client_id=None, object_id=None, resource_id=None, @@ -507,6 +540,7 @@ def _match_account(account, subscription_id, secondary_key_name, secondary_key_v set_cloud_subscription(self.cli_ctx, active_cloud.name, default_sub_id) self._storage[_SUBSCRIPTIONS] = subscriptions + return subscriptions @staticmethod def _pick_working_subscription(subscriptions): diff --git a/src/azure-cli-core/azure/cli/core/tests/test_profile.py b/src/azure-cli-core/azure/cli/core/tests/test_profile.py index 20ae4c892cb..0644fe267c2 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_profile.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_profile.py @@ -12,7 +12,8 @@ from unittest import mock from azure.cli.core._profile import (Profile, SubscriptionFinder, _attach_token_tenant, - _transform_subscription_for_multiapi) + _transform_subscription_for_multiapi, + _TENANT_LEVEL_ACCOUNT_NAME) from azure.cli.core.auth.util import AccessToken from azure.cli.core.mock import DummyCli from azure.mgmt.resource.subscriptions.models import \ @@ -551,7 +552,7 @@ def test_login_with_mi_system_assigned_no_subscriptions(self, create_subscriptio self.assertEqual(len(subscriptions), 1) s = subscriptions[0] - self.assertEqual(s['name'], 'N/A(tenant level account)') + self.assertEqual(s['name'], _TENANT_LEVEL_ACCOUNT_NAME) self.assertEqual(s['id'], self.test_mi_tenant) self.assertEqual(s['tenantId'], self.test_mi_tenant) @@ -656,7 +657,7 @@ def test_login_no_subscription(self, can_launch_browser_mock, self.assertEqual(subs[0]['id'], self.tenant_id) self.assertEqual(subs[0]['state'], 'Enabled') self.assertEqual(subs[0]['tenantId'], self.tenant_id) - self.assertEqual(subs[0]['name'], 'N/A(tenant level account)') + self.assertEqual(subs[0]['name'], _TENANT_LEVEL_ACCOUNT_NAME) self.assertTrue(profile.is_tenant_level_account()) @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) @@ -1638,6 +1639,284 @@ class SimpleManagedByTenant: assert d == {'managedByTenants': [{"tenantId": tenant_id}]} +class TestLoginSubscriptionFilter(unittest.TestCase): + """Tests for Profile.login() with --skip-subscription-discovery and --subscription parameters.""" + + @classmethod + def setUpClass(cls): + cls.tenant_id = 'test.onmicrosoft.com' + cls.user1 = 'foo@foo.com' + cls.user_identity_mock = { + 'username': cls.user1, + 'tenantId': cls.tenant_id + } + + cls.sub_id = '00000000-0000-0000-0000-000000000001' + cls.subscription_raw = SubscriptionStub( + 'subscriptions/{}'.format(cls.sub_id), + 'test sub', 'Enabled', tenant_id=cls.tenant_id, + home_tenant_id=cls.tenant_id) + + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_skip_discovery_with_subscription(self, can_launch_browser_mock, login_with_auth_code_mock, + get_user_credential_mock, create_subscription_client_mock): + """--skip-subscription-discovery --subscription S: calls GET /subscriptions/S (1 API call), S is default.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + mock_subscription_client.subscriptions.get.return_value = deepcopy(self.subscription_raw) + create_subscription_client_mock.return_value = mock_subscription_client + + storage_mock = {'subscriptions': None} + profile = Profile(cli_ctx=cli, storage=storage_mock) + subs = profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=False, + skip_subscription_discovery=True, default_subscription=self.sub_id) + + # Assert GET was called, not LIST + mock_subscription_client.subscriptions.get.assert_called_once_with(self.sub_id) + mock_subscription_client.subscriptions.list.assert_not_called() + + # Assert subscription returned and is default + self.assertEqual(len(subs), 1) + self.assertEqual(subs[0]['id'], self.sub_id) + self.assertEqual(subs[0]['name'], 'test sub') + self.assertTrue(subs[0]['isDefault']) + + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_skip_discovery_bare_mode(self, can_launch_browser_mock, login_with_auth_code_mock, + get_user_credential_mock, create_subscription_client_mock): + """--skip-subscription-discovery (no --subscription): 0 ARM calls, tenant-level account created.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + create_subscription_client_mock.return_value = mock_subscription_client + + storage_mock = {'subscriptions': None} + profile = Profile(cli_ctx=cli, storage=storage_mock) + subs = profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=False, + skip_subscription_discovery=True) + + # Assert no ARM subscription calls were made + mock_subscription_client.subscriptions.get.assert_not_called() + mock_subscription_client.subscriptions.list.assert_not_called() + + # Assert tenant-level account created + self.assertEqual(len(subs), 1) + self.assertEqual(subs[0]['id'], self.tenant_id) + self.assertEqual(subs[0]['name'], _TENANT_LEVEL_ACCOUNT_NAME) + self.assertEqual(subs[0]['tenantId'], self.tenant_id) + self.assertTrue(profile.is_tenant_level_account()) + + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_skip_discovery_with_subscription_inaccessible(self, can_launch_browser_mock, + login_with_auth_code_mock, + get_user_credential_mock, + create_subscription_client_mock): + """--skip-subscription-discovery --subscription S where S is inaccessible: raises CLIError.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + mock_subscription_client.subscriptions.get.side_effect = Exception("Not found") + create_subscription_client_mock.return_value = mock_subscription_client + + storage_mock = {'subscriptions': None} + profile = Profile(cli_ctx=cli, storage=storage_mock) + + with self.assertRaises(CLIError): + profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=False, + skip_subscription_discovery=True, default_subscription=self.sub_id) + + @mock.patch('azure.cli.core._profile.logger', autospec=True) + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_skip_discovery_with_subscription_inaccessible_allow_no_subs(self, can_launch_browser_mock, + login_with_auth_code_mock, + get_user_credential_mock, + create_subscription_client_mock, + logger_mock): + """--skip-subscription-discovery --subscription S (inaccessible) --allow-no-subscriptions: + tenant-level account created, warning logged.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + mock_subscription_client.subscriptions.get.side_effect = Exception("Not found") + create_subscription_client_mock.return_value = mock_subscription_client + + storage_mock = {'subscriptions': None} + profile = Profile(cli_ctx=cli, storage=storage_mock) + subs = profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=True, + skip_subscription_discovery=True, default_subscription=self.sub_id) + + # Tenant-level account created since subscription was inaccessible + self.assertEqual(len(subs), 1) + self.assertEqual(subs[0]['name'], _TENANT_LEVEL_ACCOUNT_NAME) + + # Warning logged about inaccessible subscription + logger_mock.warning.assert_any_call( + "Subscription '%s' not found. Profile has tenant-level account only.", + self.sub_id) + + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_no_skip_with_subscription_sets_default(self, can_launch_browser_mock, login_with_auth_code_mock, + get_user_credential_mock, + create_subscription_client_mock): + """--subscription S (without skip): full discovery unchanged, S set as default.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + mock_subscription_client.tenants.list.return_value = [TenantStub(self.tenant_id)] + mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription_raw)] + create_subscription_client_mock.return_value = mock_subscription_client + + storage_mock = {'subscriptions': None} + profile = Profile(cli_ctx=cli, storage=storage_mock) + subs = profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=False, + default_subscription=self.sub_id) + + # Assert LIST was called (full discovery), not just GET + mock_subscription_client.subscriptions.list.assert_called() + mock_subscription_client.subscriptions.get.assert_not_called() + + # Assert subscription returned and is default + self.assertEqual(len(subs), 1) + self.assertEqual(subs[0]['id'], self.sub_id) + self.assertTrue(subs[0]['isDefault']) + + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_no_skip_with_subscription_not_found_raises_error(self, can_launch_browser_mock, + login_with_auth_code_mock, + get_user_credential_mock, + create_subscription_client_mock): + """--subscription S (no skip), S not in discovered list: raises CLIError.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + mock_subscription_client.tenants.list.return_value = [TenantStub(self.tenant_id)] + mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription_raw)] + create_subscription_client_mock.return_value = mock_subscription_client + + storage_mock = {'subscriptions': None} + profile = Profile(cli_ctx=cli, storage=storage_mock) + + with self.assertRaisesRegex(CLIError, "Subscription 'non-existent' not found"): + profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=False, + default_subscription='non-existent') + + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_skip_discovery_preserves_prior_subscriptions(self, can_launch_browser_mock, + login_with_auth_code_mock, + get_user_credential_mock, + create_subscription_client_mock): + """Profile merge: prior subscriptions in azureProfile.json are preserved.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + mock_subscription_client.subscriptions.get.return_value = deepcopy(self.subscription_raw) + create_subscription_client_mock.return_value = mock_subscription_client + + # Pre-existing subscription in profile + existing_sub_id = '00000000-0000-0000-0000-000000000099' + existing_sub = { + 'id': existing_sub_id, + 'name': 'existing sub', + 'state': 'Enabled', + 'user': {'name': self.user1, 'type': 'user'}, + 'isDefault': True, + 'tenantId': self.tenant_id, + 'environmentName': 'AzureCloud' + } + storage_mock = {'subscriptions': [existing_sub]} + profile = Profile(cli_ctx=cli, storage=storage_mock) + subs = profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=False, + skip_subscription_discovery=True, default_subscription=self.sub_id) + + # Both the new and existing subscriptions should be in storage (merge) + stored = storage_mock['subscriptions'] + stored_ids = {s['id'] for s in stored} + self.assertIn(self.sub_id, stored_ids) + self.assertIn(existing_sub_id, stored_ids) + + @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential', autospec=True) + @mock.patch('azure.cli.core.auth.identity.Identity.login_with_auth_code', autospec=True) + @mock.patch('azure.cli.core._profile.can_launch_browser', autospec=True, return_value=True) + def test_skip_discovery_bare_preserves_prior_subscriptions(self, can_launch_browser_mock, + login_with_auth_code_mock, + get_user_credential_mock, + create_subscription_client_mock): + """Bare --skip-subscription-discovery preserves prior subscriptions in azureProfile.json.""" + login_with_auth_code_mock.return_value = self.user_identity_mock + + cli = DummyCli() + mock_subscription_client = mock.MagicMock() + create_subscription_client_mock.return_value = mock_subscription_client + + # Pre-existing subscription in profile + existing_sub_id = '00000000-0000-0000-0000-000000000099' + existing_sub = { + 'id': existing_sub_id, + 'name': 'existing sub', + 'state': 'Enabled', + 'user': {'name': self.user1, 'type': 'user'}, + 'isDefault': True, + 'tenantId': self.tenant_id, + 'environmentName': 'AzureCloud' + } + storage_mock = {'subscriptions': [existing_sub]} + profile = Profile(cli_ctx=cli, storage=storage_mock) + subs = profile.login(True, None, None, False, self.tenant_id, + allow_no_subscriptions=False, + skip_subscription_discovery=True) + + # No ARM calls in bare mode + mock_subscription_client.subscriptions.get.assert_not_called() + mock_subscription_client.subscriptions.list.assert_not_called() + + # Tenant-level account created + self.assertEqual(len(subs), 1) + self.assertEqual(subs[0]['name'], _TENANT_LEVEL_ACCOUNT_NAME) + + # Prior subscription preserved in storage + stored = storage_mock['subscriptions'] + stored_ids = {s['id'] for s in stored} + self.assertIn(existing_sub_id, stored_ids) + self.assertIn(self.tenant_id, stored_ids) + + class TestSubscriptionFinderFindSpecific(unittest.TestCase): """Tests for SubscriptionFinder.find_specific_subscriptions()""" diff --git a/src/azure-cli/azure/cli/command_modules/profile/__init__.py b/src/azure-cli/azure/cli/command_modules/profile/__init__.py index 8b0e3020529..859e8e5b261 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/__init__.py +++ b/src/azure-cli/azure/cli/command_modules/profile/__init__.py @@ -60,6 +60,18 @@ def load_arguments(self, command): "WWW-Authenticate header.") c.ignore('_subscription') # hide the global subscription parameter + # Skip subscription discovery + c.argument('skip_subscription_discovery', options_list=['--skip-subscription-discovery', '--skip-sub'], + action='store_true', + help='Skip the subscription discovery process during login. ' + 'Requires --tenant. Use with --subscription to ' + 'fetch a single subscription without listing all.') + c.argument('default_subscription', options_list=['--subscription', '-s'], + help='Subscription ID or name to set as the default. ' + 'When combined with --skip-subscription-discovery, ' + 'only this subscription is retrieved via a direct API call ' + '(must be a subscription ID, not a name).') + # Device code flow c.argument('use_device_code', action='store_true', help="Use device code flow. Azure CLI will also use this if it can't launch a browser, " diff --git a/src/azure-cli/azure/cli/command_modules/profile/_help.py b/src/azure-cli/azure/cli/command_modules/profile/_help.py index e61c8738c73..e13252a5a59 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/_help.py +++ b/src/azure-cli/azure/cli/command_modules/profile/_help.py @@ -47,6 +47,12 @@ text: az login --identity --client-id 00000000-0000-0000-0000-000000000000 - name: Log in with a user-assigned managed identity's resource ID. text: az login --identity --resource-id /subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/MyResourceGroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyIdentity + - name: Log in and skip subscription discovery, fetching only a single subscription (fastest for tenants with many subscriptions). + text: az login --tenant TENANT_ID --skip-subscription-discovery --subscription SUBSCRIPTION_ID + - name: Log in and skip subscription discovery entirely for tenant-level operations only (e.g., 'az ad'). + text: az login --tenant TENANT_ID --skip-subscription-discovery + - name: Log in with full discovery but set a specific subscription as the default. + text: az login --subscription SUBSCRIPTION_ID """ helps['account'] = """ diff --git a/src/azure-cli/azure/cli/command_modules/profile/custom.py b/src/azure-cli/azure/cli/command_modules/profile/custom.py index d804bb1f3d3..b1b2483f4c2 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/custom.py +++ b/src/azure-cli/azure/cli/command_modules/profile/custom.py @@ -128,7 +128,9 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_ # Service principal service_principal=None, certificate=None, use_cert_sn_issuer=None, client_assertion=None, # Managed identity - identity=False, client_id=None, object_id=None, resource_id=None): + identity=False, client_id=None, object_id=None, resource_id=None, + # Subscription discovery and default subscription selection control + skip_subscription_discovery=False, default_subscription=None): """Log in to access Azure subscriptions""" # quick argument usage check @@ -143,6 +145,8 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_ raise CLIError("usage error: '--use-sn-issuer' is only applicable with a service principal") if service_principal and not username: raise CLIError('usage error: --service-principal --username NAME --password SECRET --tenant TENANT') + if skip_subscription_discovery and not tenant: + raise CLIError("usage error: '--skip-subscription-discovery' requires '--tenant'") if username and not service_principal and not identity: if cmd.cli_ctx.cloud.endpoints.active_directory.startswith('https://login.microsoftonline.com'): logger.warning(USERNAME_PASSWORD_DEPRECATION_WARNING_AZURE_CLOUD) @@ -187,7 +191,10 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_ from azure.cli.core.telemetry import set_login_experience_v2 set_login_experience_v2(login_experience_v2) - select_subscription = interactive and sys.stdin.isatty() and sys.stdout.isatty() and login_experience_v2 + # When --subscription or --skip-subscription-discovery is provided, bypass the interactive selector + select_subscription = (interactive and sys.stdin.isatty() and sys.stdout.isatty() and + login_experience_v2 and not default_subscription and + not skip_subscription_discovery) subscriptions = profile.login( interactive, @@ -200,7 +207,9 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_ allow_no_subscriptions=allow_no_subscriptions, use_cert_sn_issuer=use_cert_sn_issuer, show_progress=select_subscription, - claims_challenge=claims_challenge + claims_challenge=claims_challenge, + skip_subscription_discovery=skip_subscription_discovery, + default_subscription=default_subscription ) # Launch interactive account selection. No JSON output. diff --git a/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py b/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py index 74547381ccf..4b5943933e6 100644 --- a/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py +++ b/src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py @@ -9,6 +9,7 @@ from azure.cli.command_modules.profile.custom import ( list_subscriptions, get_access_token, login, logout, account_clear, _remove_adal_token_cache) +from azure.cli.core._profile import _TENANT_LEVEL_ACCOUNT_NAME from azure.cli.core.mock import DummyCli from knack.util import CLIError @@ -194,3 +195,74 @@ def test_remove_adal_token_cache(self): f.write("test_token_cache") assert _remove_adal_token_cache() assert not os.path.exists(adal_token_cache) + + +class TestLoginSubscriptionFilter(unittest.TestCase): + """Tests for custom.login() with --skip-subscription-discovery and --subscription parameters.""" + + def test_skip_subscription_discovery_requires_tenant(self): + """--skip-subscription-discovery without --tenant raises CLIError.""" + cmd = mock.MagicMock() + cmd.cli_ctx = DummyCli() + with self.assertRaisesRegex(CLIError, "'--skip-subscription-discovery' requires '--tenant'"): + login(cmd, skip_subscription_discovery=True) + + def test_skip_subscription_discovery_without_tenant_with_subscription(self): + """--skip-subscription-discovery --subscription S without --tenant raises CLIError.""" + cmd = mock.MagicMock() + cmd.cli_ctx = DummyCli() + with self.assertRaisesRegex(CLIError, "'--skip-subscription-discovery' requires '--tenant'"): + login(cmd, skip_subscription_discovery=True, default_subscription='sub-id') + + @mock.patch('azure.cli.command_modules.profile._subscription_selector.SubscriptionSelector', autospec=True) + @mock.patch('azure.cli.command_modules.profile.custom.Profile', autospec=True) + def test_skip_subscription_discovery_bypasses_selector(self, profile_mock, selector_mock): + """--skip-subscription-discovery should bypass the interactive selector.""" + tenant_id = 'test-tenant' + profile_instance = mock.MagicMock() + profile_instance.login.return_value = [ + {'id': tenant_id, 'name': _TENANT_LEVEL_ACCOUNT_NAME, 'isDefault': True, + 'environmentName': 'AzureCloud', 'tenantId': tenant_id} + ] + profile_mock.return_value = profile_instance + + cmd = mock.MagicMock() + cmd.cli_ctx = DummyCli() + cmd.cli_ctx.config = mock.MagicMock() + cmd.cli_ctx.config.getboolean.return_value = True # login_experience_v2 = True + + result = login(cmd, tenant=tenant_id, skip_subscription_discovery=True) + + # Assert selector was never instantiated + selector_mock.assert_not_called() + self.assertIsNotNone(result) + assert result is not None + self.assertEqual(len(result), 1) + self.assertEqual(result[0]['name'], _TENANT_LEVEL_ACCOUNT_NAME) + + @mock.patch('azure.cli.command_modules.profile._subscription_selector.SubscriptionSelector', autospec=True) + @mock.patch('azure.cli.command_modules.profile.custom.Profile', autospec=True) + def test_default_subscription_bypasses_selector(self, profile_mock, selector_mock): + """--subscription should bypass the interactive selector even without --skip-subscription-discovery.""" + sub_id = 'target-sub-id' + profile_instance = mock.MagicMock() + profile_instance.login.return_value = [ + {'id': sub_id, 'name': 'Target Sub', 'isDefault': True, + 'environmentName': 'AzureCloud', 'tenantId': 'tenant1'}, + {'id': 'other-sub', 'name': 'Other Sub', 'isDefault': False, + 'environmentName': 'AzureCloud', 'tenantId': 'tenant1'} + ] + profile_mock.return_value = profile_instance + + cmd = mock.MagicMock() + cmd.cli_ctx = DummyCli() + cmd.cli_ctx.config = mock.MagicMock() + cmd.cli_ctx.config.getboolean.return_value = False + + # Interactive login (no username) with --subscription + result = login(cmd, tenant='tenant1', default_subscription=sub_id) + + # Assert selector was never instantiated + selector_mock.assert_not_called() + assert result is not None + self.assertEqual(len(result), 2)