Fixes how azure auth is handled on the azure pane (#9654)

extensions/azurecore/src/azureResource/commands.ts

@@ -76,7 +76,7 @@ export function registerAzureResourceCommands(appContext: AppContext, tree: Azur
 				const token = tokens[tenant.id].token;
 				const tokenType = tokens[tenant.id].tokenType;
 
-				result.resourceGroups.push(...await service.getResources(subscription, new TokenCredentials(token, tokenType)));
+				result.resourceGroups.push(...await service.getResources(subscription, new TokenCredentials(token, tokenType), account));
 			} catch (err) {
 				const error = new Error(localize('azure.accounts.getResourceGroups.queryError', "Error fetching resource groups for account {0} ({1}) subscription {2} ({3}) tenant {4} : {5}",
 					account.displayInfo.displayName,

extensions/azurecore/src/azureResource/interfaces.ts

@@ -33,7 +33,7 @@ export interface IAzureResourceCacheService {
 }
 
 export interface IAzureResourceTenantService {
-	getTenantId(subscription: azureResource.AzureResourceSubscription): Promise<string>;
+	getTenantId(subscription: azureResource.AzureResourceSubscription, account: Account, credential: msRest.ServiceClientCredentials): Promise<string>;
 }
 
 export interface IAzureResourceNodeWithProviderId {
@@ -42,5 +42,5 @@ export interface IAzureResourceNodeWithProviderId {
 }
 
 export interface IAzureResourceService<T extends azureResource.AzureResource> {
-	getResources(subscription: azureResource.AzureResourceSubscription, credential: msRest.ServiceClientCredentials): Promise<T[]>;
+	getResources(subscription: azureResource.AzureResourceSubscription, credential: msRest.ServiceClientCredentials, account: Account): Promise<T[]>;
 }

extensions/azurecore/src/azureResource/providers/database/databaseService.ts

@@ -9,14 +9,15 @@ import { IAzureResourceService } from '../../interfaces';
 import { serversQuery, DbServerGraphData } from '../databaseServer/databaseServerService';
 import { ResourceGraphClient } from '@azure/arm-resourcegraph';
 import { queryGraphResources, GraphData } from '../resourceTreeDataProviderBase';
+import { Account } from 'azdata';
 
 interface DatabaseGraphData extends GraphData {
 	kind: string;
 }
 export class AzureResourceDatabaseService implements IAzureResourceService<azureResource.AzureResourceDatabase> {
-	public async getResources(subscription: azureResource.AzureResourceSubscription, credential: ServiceClientCredentials): Promise<azureResource.AzureResourceDatabase[]> {
+	public async getResources(subscription: azureResource.AzureResourceSubscription, credential: ServiceClientCredentials, account: Account): Promise<azureResource.AzureResourceDatabase[]> {
 		const databases: azureResource.AzureResourceDatabase[] = [];
-		const resourceClient = new ResourceGraphClient(credential);
+		const resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
 
 		// Query servers and databases in parallel (start both promises before waiting on the 1st)
 		let serverQueryPromise = queryGraphResources<GraphData>(resourceClient, subscription.id, serversQuery);

extensions/azurecore/src/azureResource/providers/database/databaseTreeDataProvider.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { TreeItem, ExtensionNodeType } from 'azdata';
+import { TreeItem, ExtensionNodeType, Account } from 'azdata';
 import { TreeItemCollapsibleState, ExtensionContext } from 'vscode';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
@@ -27,7 +27,7 @@ export class AzureResourceDatabaseTreeDataProvider extends ResourceTreeDataProvi
 	) {
 		super(databaseService, apiWrapper);
 	}
-	protected getTreeItemForResource(database: azureResource.AzureResourceDatabase): TreeItem {
+	protected getTreeItemForResource(database: azureResource.AzureResourceDatabase, account: Account): TreeItem {
 		return {
 			id: `databaseServer_${database.serverFullName}.database_${database.name}`,
 			label: `${database.name} (${database.serverName})`,
@@ -44,13 +44,14 @@ export class AzureResourceDatabaseTreeDataProvider extends ResourceTreeDataProvi
 				databaseName: database.name,
 				userName: database.loginName,
 				password: '',
-				authenticationType: 'SqlLogin',
+				authenticationType: 'AzureMFA',
 				savePassword: true,
 				groupFullName: '',
 				groupId: '',
 				providerName: 'MSSQL',
 				saveProfile: false,
-				options: {}
+				options: {},
+				azureAccount: account.key.accountId
 			},
 			childProvider: 'MSSQL',
 			type: ExtensionNodeType.Database

extensions/azurecore/src/azureResource/providers/databaseServer/databaseServerTreeDataProvider.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ExtensionNodeType, TreeItem } from 'azdata';
+import { ExtensionNodeType, TreeItem, Account } from 'azdata';
 import { TreeItemCollapsibleState, ExtensionContext } from 'vscode';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
@@ -28,7 +28,7 @@ export class AzureResourceDatabaseServerTreeDataProvider extends ResourceTreeDat
 	}
 
 
-	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer): TreeItem {
+	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer, account: Account): TreeItem {
 		return {
 			id: `databaseServer_${databaseServer.id ? databaseServer.id : databaseServer.name}`,
 			label: databaseServer.name,
@@ -43,15 +43,16 @@ export class AzureResourceDatabaseServerTreeDataProvider extends ResourceTreeDat
 				connectionName: undefined,
 				serverName: databaseServer.fullName,
 				databaseName: databaseServer.defaultDatabaseName,
-				userName: databaseServer.loginName,
+				userName: '',
 				password: '',
-				authenticationType: 'SqlLogin',
+				authenticationType: 'AzureMFA',
 				savePassword: true,
 				groupFullName: '',
 				groupId: '',
 				providerName: 'MSSQL',
 				saveProfile: false,
-				options: {}
+				options: {},
+				azureAccount: account.key.accountId
 			},
 			childProvider: 'MSSQL',
 			type: ExtensionNodeType.Server

extensions/azurecore/src/azureResource/providers/postgresArcServer/postgresServerTreeDataProvider.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ExtensionNodeType, TreeItem } from 'azdata';
+import { ExtensionNodeType, TreeItem, Account } from 'azdata';
 import { TreeItemCollapsibleState, ExtensionContext } from 'vscode';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
@@ -28,7 +28,7 @@ export class PostgresServerArcTreeDataProvider extends ResourceTreeDataProviderB
 	}
 
 
-	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer): TreeItem {
+	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer, account: Account): TreeItem {
 		return {
 			id: `databaseServer_${databaseServer.id ? databaseServer.id : databaseServer.name}`,
 			label: databaseServer.name,
@@ -45,7 +45,7 @@ export class PostgresServerArcTreeDataProvider extends ResourceTreeDataProviderB
 				databaseName: databaseServer.defaultDatabaseName,
 				userName: `${databaseServer.loginName}@${databaseServer.fullName}`,
 				password: '',
-				authenticationType: 'SqlLogin',
+				authenticationType: 'AzureMFA',
 				savePassword: true,
 				groupFullName: '',
 				groupId: '',
@@ -54,7 +54,8 @@ export class PostgresServerArcTreeDataProvider extends ResourceTreeDataProviderB
 				options: {
 					// Set default for SSL or will get error complaining about it not being set correctly
 					'sslmode': 'require'
-				}
+				},
+				azureAccount: account.key.accountId
 			},
 			childProvider: 'PGSQL',
 			type: ExtensionNodeType.Server

extensions/azurecore/src/azureResource/providers/postgresServer/postgresServerTreeDataProvider.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ExtensionNodeType, TreeItem } from 'azdata';
+import { ExtensionNodeType, TreeItem, Account } from 'azdata';
 import { TreeItemCollapsibleState, ExtensionContext } from 'vscode';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
@@ -28,7 +28,7 @@ export class PostgresServerTreeDataProvider extends ResourceTreeDataProviderBase
 	}
 
 
-	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer): TreeItem {
+	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer, account: Account): TreeItem {
 		return {
 			id: `databaseServer_${databaseServer.id ? databaseServer.id : databaseServer.name}`,
 			label: databaseServer.name,
@@ -45,7 +45,7 @@ export class PostgresServerTreeDataProvider extends ResourceTreeDataProviderBase
 				databaseName: databaseServer.defaultDatabaseName,
 				userName: `${databaseServer.loginName}@${databaseServer.fullName}`,
 				password: '',
-				authenticationType: 'SqlLogin',
+				authenticationType: 'AzureMFA',
 				savePassword: true,
 				groupFullName: '',
 				groupId: '',
@@ -54,7 +54,8 @@ export class PostgresServerTreeDataProvider extends ResourceTreeDataProviderBase
 				options: {
 					// Set default for SSL or will get error complaining about it not being set correctly
 					'sslmode': 'require'
-				}
+				},
+				azureAccount: account.key.accountId
 			},
 			childProvider: 'PGSQL',
 			type: ExtensionNodeType.Server

extensions/azurecore/src/azureResource/providers/resourceTreeDataProviderBase.ts

@@ -36,7 +36,7 @@ export abstract class ResourceTreeDataProviderBase<T extends azureResource.Azure
 				account: element.account,
 				subscription: element.subscription,
 				tenantId: element.tenantId,
-				treeItem: this.getTreeItemForResource(resource)
+				treeItem: this.getTreeItemForResource(resource, element.account)
 			}).sort((a, b) => a.treeItem.label.localeCompare(b.treeItem.label));
 		} catch (error) {
 			console.log(AzureResourceErrorMessageUtil.getErrorMessage(error));
@@ -48,11 +48,11 @@ export abstract class ResourceTreeDataProviderBase<T extends azureResource.Azure
 		const tokens = await this._apiWrapper.getSecurityToken(element.account, azdata.AzureResource.ResourceManagement);
 		const credential = new msRest.TokenCredentials(tokens[element.tenantId].token, tokens[element.tenantId].tokenType);
 
-		const resources: T[] = await this._resourceService.getResources(element.subscription, credential) || <T[]>[];
+		const resources: T[] = await this._resourceService.getResources(element.subscription, credential, element.account) || <T[]>[];
 		return resources;
 	}
 
-	protected abstract getTreeItemForResource(resource: T): azdata.TreeItem;
+	protected abstract getTreeItemForResource(resource: T, account: azdata.Account): azdata.TreeItem;
 
 	protected abstract createContainerNode(): azureResource.IAzureResourceNode;
 }
@@ -121,9 +121,9 @@ export abstract class ResourceServiceBase<T extends GraphData, U extends azureRe
 	 */
 	protected abstract get query(): string;
 
-	public async getResources(subscription: azureResource.AzureResourceSubscription, credential: msRest.ServiceClientCredentials): Promise<U[]> {
+	public async getResources(subscription: azureResource.AzureResourceSubscription, credential: msRest.ServiceClientCredentials, account: azdata.Account): Promise<U[]> {
 		const convertedResources: U[] = [];
-		const resourceClient = new ResourceGraphClient(credential);
+		const resourceClient = new ResourceGraphClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
 		let graphResources = await queryGraphResources<T>(resourceClient, subscription.id, this.query);
 		let ids = new Set<string>();
 		graphResources.forEach((res) => {

extensions/azurecore/src/azureResource/providers/sqlinstance/sqlInstanceTreeDataProvider.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ExtensionNodeType, TreeItem } from 'azdata';
+import { ExtensionNodeType, TreeItem, Account } from 'azdata';
 import { TreeItemCollapsibleState, ExtensionContext } from 'vscode';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
@@ -28,7 +28,7 @@ export class SqlInstanceTreeDataProvider extends ResourceTreeDataProviderBase<az
 	}
 
 
-	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer): TreeItem {
+	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer, account: Account): TreeItem {
 		return {
 			id: `sqlInstance_${databaseServer.id ? databaseServer.id : databaseServer.name}`,
 			label: databaseServer.name,
@@ -45,13 +45,14 @@ export class SqlInstanceTreeDataProvider extends ResourceTreeDataProviderBase<az
 				databaseName: databaseServer.defaultDatabaseName,
 				userName: databaseServer.loginName,
 				password: '',
-				authenticationType: 'SqlLogin',
+				authenticationType: 'AzureMFA',
 				savePassword: true,
 				groupFullName: '',
 				groupId: '',
 				providerName: 'MSSQL',
 				saveProfile: false,
-				options: {}
+				options: {},
+				azureAccount: account.key.accountId
 			},
 			childProvider: 'MSSQL',
 			type: ExtensionNodeType.Server

extensions/azurecore/src/azureResource/providers/sqlinstanceArc/sqlInstanceArcTreeDataProvider.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ExtensionNodeType, TreeItem } from 'azdata';
+import { ExtensionNodeType, TreeItem, Account } from 'azdata';
 import { TreeItemCollapsibleState, ExtensionContext } from 'vscode';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
@@ -28,7 +28,7 @@ export class SqlInstanceArcTreeDataProvider extends ResourceTreeDataProviderBase
 	}
 
 
-	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer): TreeItem {
+	protected getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer, account: Account): TreeItem {
 		return {
 			id: `sqlInstance_${databaseServer.id ? databaseServer.id : databaseServer.name}`,
 			label: databaseServer.name,
@@ -45,13 +45,14 @@ export class SqlInstanceArcTreeDataProvider extends ResourceTreeDataProviderBase
 				databaseName: databaseServer.defaultDatabaseName,
 				userName: databaseServer.loginName,
 				password: '',
-				authenticationType: 'SqlLogin',
+				authenticationType: 'AzureMFA',
 				savePassword: true,
 				groupFullName: '',
 				groupId: '',
 				providerName: 'MSSQL',
 				saveProfile: false,
-				options: {}
+				options: {},
+				azureAccount: account.key.accountId
 			},
 			childProvider: 'MSSQL',
 			type: ExtensionNodeType.Server

extensions/azurecore/src/azureResource/services/subscriptionService.ts

@@ -4,17 +4,16 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Account } from 'azdata';
-import { ServiceClientCredentials } from '@azure/ms-rest-js';
 import { SubscriptionClient } from '@azure/arm-subscriptions';
 
 import { azureResource } from '../azure-resource';
 import { IAzureResourceSubscriptionService } from '../interfaces';
 
 export class AzureResourceSubscriptionService implements IAzureResourceSubscriptionService {
-	public async getSubscriptions(account: Account, credential: ServiceClientCredentials): Promise<azureResource.AzureResourceSubscription[]> {
+	public async getSubscriptions(account: Account, credential: any): Promise<azureResource.AzureResourceSubscription[]> {
 		const subscriptions: azureResource.AzureResourceSubscription[] = [];
 
-		const subClient = new SubscriptionClient(credential);
+		const subClient = new SubscriptionClient(credential, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
 		const subs = await subClient.subscriptions.list();
 		subs.forEach((sub) => subscriptions.push({
 			id: sub.subscriptionId,

extensions/azurecore/src/azureResource/services/tenantService.ts

@@ -2,29 +2,17 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-
-import * as request from 'request';
+import { SubscriptionClient } from '@azure/arm-subscriptions';
 
 import { azureResource } from '../azure-resource';
 import { IAzureResourceTenantService } from '../interfaces';
+import { Account } from 'azdata';
 
 export class AzureResourceTenantService implements IAzureResourceTenantService {
-	public async getTenantId(subscription: azureResource.AzureResourceSubscription): Promise<string> {
-		const requestPromisified = new Promise<string>((resolve, reject) => {
-			const url = `https://management.azure.com/subscriptions/${subscription.id}?api-version=2014-04-01`;
-			request(url, function (error, response, body) {
-				if (response.statusCode === 401) {
-					const tenantIdRegEx = /authorization_uri="https:\/\/login\.windows\.net\/([0-9a-fA-F]{8}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{12})"/;
-					const teantIdString = response.headers['www-authenticate'];
-					if (tenantIdRegEx.test(teantIdString)) {
-						resolve(tenantIdRegEx.exec(teantIdString)[1]);
-					} else {
-						reject();
-					}
-				}
-			});
-		});
+	public async getTenantId(subscription: azureResource.AzureResourceSubscription, account: Account, credentials: any): Promise<string> {
+		const subClient = new SubscriptionClient(credentials, { baseUri: account.properties.providerSettings.settings.armResource.endpoint });
 
-		return await requestPromisified;
+		const result = await subClient.subscriptions.get(subscription.id);
+		return result.subscriptionId;
 	}
 }

extensions/azurecore/src/azureResource/tree/accountTreeNode.ts

@@ -42,11 +42,10 @@ export class AzureResourceAccountTreeNode extends AzureResourceContainerTreeNode
 	public async getChildren(): Promise<TreeNode[]> {
 		try {
 			let subscriptions: azureResource.AzureResourceSubscription[] = [];
+			const tokens = await this.appContext.apiWrapper.getSecurityToken(this.account, AzureResource.ResourceManagement);
 
 			if (this._isClearingCache) {
 				try {
-					const tokens = await this.appContext.apiWrapper.getSecurityToken(this.account, AzureResource.ResourceManagement);
-
 					for (const tenant of this.account.properties.tenants) {
 						const token = tokens[tenant.id].token;
 						const tokenType = tokens[tenant.id].tokenType;
@@ -56,7 +55,6 @@ export class AzureResourceAccountTreeNode extends AzureResourceContainerTreeNode
 				} catch (error) {
 					throw new AzureResourceCredentialError(localize('azure.resource.tree.accountTreeNode.credentialError', "Failed to get credential for account {0}. Please refresh the account.", this.account.key.accountId), error);
 				}
-
 				this.updateCache<azureResource.AzureResourceSubscription[]>(subscriptions);
 
 				this._isClearingCache = false;
@@ -82,7 +80,8 @@ export class AzureResourceAccountTreeNode extends AzureResourceContainerTreeNode
 				return [AzureResourceMessageTreeNode.create(AzureResourceAccountTreeNode.noSubscriptionsLabel, this)];
 			} else {
 				let subTreeNodes = await Promise.all(subscriptions.map(async (subscription) => {
-					const tenantId = await this._tenantService.getTenantId(subscription);
+					const token = tokens[subscription.id];
+					const tenantId = await this._tenantService.getTenantId(subscription, this.account, new TokenCredentials(token.token, token.tokenType));
 
 					return new AzureResourceSubscriptionTreeNode(this.account, subscription, tenantId, this.appContext, this.treeChangeHandler, this);
 				}));

extensions/azurecore/src/azureResource/tree/treeProvider.ts

@@ -37,13 +37,16 @@ export class AzureResourceTreeProvider implements TreeDataProvider<TreeNode>, IA
 	private hookAccountService(appContext: AppContext): void {
 		this.accountService = appContext.getService<IAzureResourceAccountService>(AzureResourceServiceNames.accountService);
 		if (this.accountService) {
-			this.accountService.onDidChangeAccounts((e: azdata.DidChangeAccountsParams) => {
+			this.accountService.onDidChangeAccounts(async (e: azdata.DidChangeAccountsParams) => {
+				// This event sends it per provider, we need to make sure we get all the azure related accounts
+				let accounts = await this.accountService.getAccounts();
+				accounts = accounts.filter(a => a.key.providerId.startsWith('azure'));
 				// the onDidChangeAccounts event will trigger in many cases where the accounts didn't actually change
 				// the notifyNodeChanged event triggers a refresh which triggers a getChildren which can trigger this callback
 				// this below check short-circuits the infinite callback loop
 				this.setSystemInitialized();
-				if (!equals(e.accounts, this.accounts)) {
-					this.accounts = e.accounts;
+				if (!equals(accounts, this.accounts)) {
+					this.accounts = accounts;
 					this.notifyNodeChanged(undefined);
 				}
 			});

extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts

@@ -104,7 +104,7 @@ describe('AzureResourceDatabaseTreeDataProvider.getChildren', function (): void
 		mockExtensionContext = TypeMoq.Mock.ofType<vscode.ExtensionContext>();
 
 		mockApiWrapper.setup((o) => o.getSecurityToken(mockAccount, azdata.AzureResource.ResourceManagement)).returns(() => Promise.resolve(mockTokens));
-		mockDatabaseService.setup((o) => o.getResources(mockSubscription, TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabases));
+		mockDatabaseService.setup((o) => o.getResources(mockSubscription, TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabases));
 		mockExtensionContext.setup((o) => o.asAbsolutePath(TypeMoq.It.isAnyString())).returns(() => TypeMoq.It.isAnyString());
 	});
 

extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts

@@ -104,7 +104,7 @@ describe('AzureResourceDatabaseServerTreeDataProvider.getChildren', function ():
 		mockExtensionContext = TypeMoq.Mock.ofType<vscode.ExtensionContext>();
 
 		mockApiWrapper.setup((o) => o.getSecurityToken(mockAccount, azdata.AzureResource.ResourceManagement)).returns(() => Promise.resolve(mockTokens));
-		mockDatabaseServerService.setup((o) => o.getResources(mockSubscription, TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabaseServers));
+		mockDatabaseServerService.setup((o) => o.getResources(mockSubscription, TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockDatabaseServers));
 		mockExtensionContext.setup((o) => o.asAbsolutePath(TypeMoq.It.isAnyString())).returns(() => TypeMoq.It.isAnyString());
 	});
 

extensions/azurecore/src/test/azureResource/tree/accountTreeNode.test.ts

@@ -76,10 +76,13 @@ const mockSubscriptions = [mockSubscription1, mockSubscription2];
 const mockFilteredSubscriptions = [mockSubscription1];
 
 const mockTokens: { [key: string]: any } = {};
-mockTokens[mockTenantId] = {
-	token: 'mock_token',
-	tokenType: 'Bearer'
-};
+
+[mockSubscription1.id, mockSubscription2.id, mockTenantId].forEach(s => {
+	mockTokens[s] = {
+		token: 'mock_token',
+		tokenType: 'Bearer'
+	};
+});
 
 const mockCredential = new TokenCredentials(mockTokens[mockTenantId].token, mockTokens[mockTenantId].tokenType);
 
@@ -108,7 +111,7 @@ describe('AzureResourceAccountTreeNode.info', function (): void {
 		mockCacheService.setup((o) => o.generateKey(TypeMoq.It.isAnyString())).returns(() => generateGuid());
 		mockCacheService.setup((o) => o.get(TypeMoq.It.isAnyString())).returns(() => mockSubscriptionCache);
 		mockCacheService.setup((o) => o.update(TypeMoq.It.isAnyString(), TypeMoq.It.isAny())).returns(() => mockSubscriptionCache = mockSubscriptions);
-		mockTenantService.setup((o) => o.getTenantId(TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
+		mockTenantService.setup((o) => o.getTenantId(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
 	});
 
 	it('Should be correct when created.', async function (): Promise<void> {
@@ -196,7 +199,7 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 		mockCacheService.setup((o) => o.generateKey(TypeMoq.It.isAnyString())).returns(() => generateGuid());
 		mockCacheService.setup((o) => o.get(TypeMoq.It.isAnyString())).returns(() => mockSubscriptionCache);
 		mockCacheService.setup((o) => o.update(TypeMoq.It.isAnyString(), TypeMoq.It.isAny())).returns(() => mockSubscriptionCache = mockSubscriptions);
-		mockTenantService.setup((o) => o.getTenantId(TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
+		mockTenantService.setup((o) => o.getTenantId(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
 	});
 
 	it('Should load subscriptions from scratch and update cache when it is clearing cache.', async function (): Promise<void> {
@@ -207,7 +210,6 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 
 		const children = await accountTreeNode.getChildren();
 
-		mockApiWrapper.verify((o) => o.getSecurityToken(mockAccount, azdata.AzureResource.ResourceManagement), TypeMoq.Times.once());
 		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential), TypeMoq.Times.once());
 		mockCacheService.verify((o) => o.get(TypeMoq.It.isAnyString()), TypeMoq.Times.exactly(0));
 		mockCacheService.verify((o) => o.update(TypeMoq.It.isAnyString(), TypeMoq.It.isAny()), TypeMoq.Times.once());
@@ -242,7 +244,7 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 		await accountTreeNode.getChildren();
 		const children = await accountTreeNode.getChildren();
 
-		mockApiWrapper.verify((o) => o.getSecurityToken(mockAccount, azdata.AzureResource.ResourceManagement), TypeMoq.Times.once());
+
 		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential), TypeMoq.Times.once());
 		mockCacheService.verify((o) => o.get(TypeMoq.It.isAnyString()), TypeMoq.Times.once());
 		mockCacheService.verify((o) => o.update(TypeMoq.It.isAnyString(), TypeMoq.It.isAny()), TypeMoq.Times.once());
@@ -331,11 +333,11 @@ describe('AzureResourceAccountTreeNode.clearCache', function (): void {
 		mockAppContext.registerService<IAzureResourceSubscriptionFilterService>(AzureResourceServiceNames.subscriptionFilterService, mockSubscriptionFilterService.object);
 		mockAppContext.registerService<IAzureResourceTenantService>(AzureResourceServiceNames.tenantService, mockTenantService.object);
 
-		mockApiWrapper.setup((o) => o.getSecurityToken(mockAccount, azdata.AzureResource.ResourceManagement)).returns(() => Promise.resolve(mockTokens));
+		mockApiWrapper.setup((o,) => o.getSecurityToken(mockAccount, azdata.AzureResource.ResourceManagement)).returns(() => Promise.resolve(mockTokens));
 		mockCacheService.setup((o) => o.generateKey(TypeMoq.It.isAnyString())).returns(() => generateGuid());
 		mockCacheService.setup((o) => o.get(TypeMoq.It.isAnyString())).returns(() => mockSubscriptionCache);
 		mockCacheService.setup((o) => o.update(TypeMoq.It.isAnyString(), TypeMoq.It.isAny())).returns(() => mockSubscriptionCache = mockSubscriptions);
-		mockTenantService.setup((o) => o.getTenantId(TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
+		mockTenantService.setup((o) => o.getTenantId(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
 	});
 
 	it('Should clear cache.', async function (): Promise<void> {