Large cleanup of AzureCore - Introduction of getAccountSecurityToken and deprecation of getSecurityToken (#11446)

* do a large cleanup of azurecore * Fix tests * Rework Device Code * Fix tests * Fix AE scenario * Fix firewall rule - clenaup logging * Shorthand syntax * Fix firewall tests * Start on tests for azureAuth * Add more tests * Address comments * Add a few more important tests * Don't throw error on old code * Fill in todo
2026-02-16 10:58:30 -05:00 · 2020-07-22 15:03:42 -07:00
parent a61b85c9ff
commit 587abd43c2
40 changed files with 1045 additions and 895 deletions
--- a/extensions/azurecore/src/test/account-provider/auths/azureAuth.test.ts
+++ b/extensions/azurecore/src/test/account-provider/auths/azureAuth.test.ts
@@ -4,114 +4,240 @@
 *--------------------------------------------------------------------------------------------*/

 import * as should from 'should';
-import * as os from 'os';
+import * as TypeMoq from 'typemoq';
+// import * as azdata from 'azdata';
+// import * as vscode from 'vscode';
+// import * as sinon from 'sinon';
 import 'mocha';
-import * as vscode from 'vscode';
-
-import { PromptFailedResult, AccountKey } from 'azdata';
-import { AzureAuth, AccessToken, RefreshToken, TokenClaims, TokenRefreshResponse } from '../../../account-provider/auths/azureAuth';
-import { AzureAccount, AzureAuthType, Deferred, Tenant } from '../../../account-provider/interfaces';
+import { AzureAuthCodeGrant } from '../../../account-provider/auths/azureAuthCodeGrant';
+// import { AzureDeviceCode } from '../../../account-provider/auths/azureDeviceCode';
+import { Token, TokenClaims, AccessToken, RefreshToken, OAuthTokenResponse, TokenPostData } from '../../../account-provider/auths/azureAuth';
+import { Tenant, AzureAccount } from '../../../account-provider/interfaces';
 import providerSettings from '../../../account-provider/providerSettings';
-import { SimpleTokenCache } from '../../../account-provider/simpleTokenCache';
-import { CredentialsTestProvider } from '../../stubs/credentialsTestProvider';
+import { AzureResource } from 'azdata';
+import { AxiosResponse } from 'axios';

-class BasicAzureAuth extends AzureAuth {
-	public async login(): Promise<AzureAccount | PromptFailedResult> {
-		throw new Error('Method not implemented.');
-	}

-	public async autoOAuthCancelled(): Promise<void> {
-		throw new Error('Method not implemented.');
-	}
+let azureAuthCodeGrant: TypeMoq.IMock<AzureAuthCodeGrant>;
+// let azureDeviceCode: TypeMoq.IMock<AzureDeviceCode>;

-	public async promptForConsent(): Promise<{ tokenRefreshResponse: TokenRefreshResponse, authCompleteDeferred: Deferred<void> } | undefined> {
-		throw new Error('Method not implemented.');
-	}
-}
+const mockToken: Token = {
+	key: 'someUniqueId',
+	token: 'test_token',
+	tokenType: 'Bearer'
+};
+let mockAccessToken: AccessToken;
+let mockRefreshToken: RefreshToken;

-let baseAuth: AzureAuth;
+const mockClaims = {
+	name: 'Name',
+	email: 'example@example.com',
+	sub: 'someUniqueId'
+} as TokenClaims;

-const accountKey: AccountKey = {
-	accountId: 'SomeAccountKey',
-	providerId: 'providerId',
+const mockTenant: Tenant = {
+	displayName: 'Tenant Name',
+	id: 'tenantID',
+	tenantCategory: 'Home',
+	userId: 'test_user'
 };

-const accessToken: AccessToken = {
-	key: accountKey.accountId,
-	token: '123'
-};
+let mockAccount: AzureAccount;

-const refreshToken: RefreshToken = {
-	key: accountKey.accountId,
-	token: '321'
-};
+const provider = providerSettings[0].metadata;

-const resourceId = 'resource';
-const tenantId = 'tenant';
+describe('Azure Authentication', function () {
+	beforeEach(function () {
+		azureAuthCodeGrant = TypeMoq.Mock.ofType<AzureAuthCodeGrant>(AzureAuthCodeGrant, TypeMoq.MockBehavior.Loose, true, provider);
+		// azureDeviceCode = TypeMoq.Mock.ofType<AzureDeviceCode>();

-const tenant: Tenant = {
-	id: tenantId,
-	displayName: 'common'
-};
+		azureAuthCodeGrant.callBase = true;
+		// authDeviceCode.callBase = true;

-// These tests don't work on Linux systems because gnome-keyring doesn't like running on headless machines.
-describe('AccountProvider.AzureAuth', function (): void {
-	beforeEach(async function (): Promise<void> {
-		const tokenCache = new SimpleTokenCache('testTokenService', os.tmpdir(), true, new CredentialsTestProvider());
-		await tokenCache.init();
-		baseAuth = new BasicAzureAuth(providerSettings[0].metadata, tokenCache, undefined, undefined, AzureAuthType.AuthCodeGrant, 'Auth Code Grant');
+		mockAccount = {
+			isStale: false,
+			properties: {
+				tenants: [mockTenant]
+			}
+		} as AzureAccount;
+
+		mockAccessToken = {
+			...mockToken
+		};
+		mockRefreshToken = {
+			...mockToken
+		};
 	});

-	it('Basic token set and get', async function (): Promise<void> {
-		await baseAuth.setCachedToken(accountKey, accessToken, refreshToken);
-		const result = await baseAuth.getCachedToken(accountKey);
+	it('accountHydration should yield a valid account', async function () {

-		should(JSON.stringify(result.accessToken)).be.equal(JSON.stringify(accessToken));
-		should(JSON.stringify(result.refreshToken)).be.equal(JSON.stringify(refreshToken));
+		azureAuthCodeGrant.setup(x => x.getTenants(mockToken)).returns((): Promise<Tenant[]> => {
+			return Promise.resolve([
+				mockTenant
+			]);
+		});
+
+		const response = await azureAuthCodeGrant.object.hydrateAccount(mockToken, mockClaims);
+		should(response.displayInfo.displayName).be.equal(`${mockClaims.name} - ${mockClaims.email}`, 'Account name should match');
+		should(response.displayInfo.userId).be.equal(mockClaims.sub, 'Account ID should match');
+		should(response.properties.tenants).be.deepEqual([mockTenant], 'Tenants should match');
 	});

-	it('Token set and get with tenant and resource id', async function (): Promise<void> {
-		await baseAuth.setCachedToken(accountKey, accessToken, refreshToken, resourceId, tenantId);
-		let result = await baseAuth.getCachedToken(accountKey, resourceId, tenantId);
+	describe('getAccountSecurityToken', function () {
+		it('should be undefined on stale account', async function () {
+			mockAccount.isStale = true;
+			const securityToken = await azureAuthCodeGrant.object.getAccountSecurityToken(mockAccount, TypeMoq.It.isAny(), TypeMoq.It.isAny());
+			should(securityToken).be.undefined();
+		});
+		it('dont find correct resources', async function () {
+			const securityToken = await azureAuthCodeGrant.object.getAccountSecurityToken(mockAccount, TypeMoq.It.isAny(), -1);
+			should(securityToken).be.undefined();
+		});
+		it('incorrect tenant', async function () {
+			await azureAuthCodeGrant.object.getAccountSecurityToken(mockAccount, 'invalid_tenant', AzureResource.MicrosoftResourceManagement).should.be.rejected();
+		});

-		should(JSON.stringify(result.accessToken)).be.equal(JSON.stringify(accessToken));
-		should(JSON.stringify(result.refreshToken)).be.equal(JSON.stringify(refreshToken));
+		it('saved token exists and can be reused', async function () {
+			delete (mockAccessToken as any).tokenType;
+			azureAuthCodeGrant.setup(x => x.getSavedToken(mockTenant, provider.settings.microsoftResource, mockAccount.key)).returns((): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken, expiresOn: string }> => {
+				return Promise.resolve({
+					accessToken: mockAccessToken,
+					refreshToken: mockRefreshToken,
+					expiresOn: `${(new Date().getTime() / 1000) + (10 * 60)}`
+				})
+			});
+			const securityToken = await azureAuthCodeGrant.object.getAccountSecurityToken(mockAccount, mockTenant.id, AzureResource.MicrosoftResourceManagement);
+
+			should(securityToken.tokenType).be.equal('Bearer', 'tokenType should be bearer on a successful getSecurityToken from cache')
+		});
+
+
+		it('saved token had invalid expiration', async function () {
+			delete (mockAccessToken as any).tokenType;
+			(mockAccessToken as any).invalidData = 'this should not exist on response';
+			azureAuthCodeGrant.setup(x => x.getSavedToken(mockTenant, provider.settings.microsoftResource, mockAccount.key)).returns((): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken, expiresOn: string }> => {
+				return Promise.resolve({
+					accessToken: mockAccessToken,
+					refreshToken: mockRefreshToken,
+					expiresOn: undefined
+				});
+			});
+			azureAuthCodeGrant.setup(x => x.refreshToken(mockTenant, provider.settings.microsoftResource, mockRefreshToken)).returns((): Promise<OAuthTokenResponse> => {
+				const mockToken: AccessToken = JSON.parse(JSON.stringify(mockAccessToken));
+				delete (mockToken as any).invalidData;
+				return Promise.resolve({
+					accessToken: mockToken
+				} as OAuthTokenResponse);
+			});
+			const securityToken = await azureAuthCodeGrant.object.getAccountSecurityToken(mockAccount, mockTenant.id, AzureResource.MicrosoftResourceManagement);
+
+			should((securityToken as any).invalidData).be.undefined(); // Ensure its a new one
+			should(securityToken.tokenType).be.equal('Bearer', 'tokenType should be bearer on a successful getSecurityToken from cache')
+
+			azureAuthCodeGrant.verify(x => x.refreshToken(mockTenant, provider.settings.microsoftResource, mockRefreshToken), TypeMoq.Times.once());
+		});
+
+		describe('no saved token', function () {
+			it('no base token', async function () {
+				azureAuthCodeGrant.setup(x => x.getSavedToken(mockTenant, provider.settings.microsoftResource, mockAccount.key)).returns((): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken, expiresOn: string }> => {
+					return Promise.resolve(undefined);
+				});
+
+				azureAuthCodeGrant.setup(x => x.getSavedToken(azureAuthCodeGrant.object.commonTenant, provider.settings.microsoftResource, mockAccount.key)).returns((): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken, expiresOn: string }> => {
+					return Promise.resolve(undefined);
+				});
+
+				await azureAuthCodeGrant.object.getAccountSecurityToken(mockAccount, mockTenant.id, AzureResource.MicrosoftResourceManagement).should.be.rejected();
+			});
+
+			it('base token exists', async function () {
+				azureAuthCodeGrant.setup(x => x.getSavedToken(mockTenant, provider.settings.microsoftResource, mockAccount.key)).returns((): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken, expiresOn: string }> => {
+					return Promise.resolve(undefined);
+				});
+
+				azureAuthCodeGrant.setup(x => x.getSavedToken(azureAuthCodeGrant.object.commonTenant, provider.settings.microsoftResource, mockAccount.key)).returns((): Promise<{ accessToken: AccessToken, refreshToken: RefreshToken, expiresOn: string }> => {
+					return Promise.resolve({
+						accessToken: mockAccessToken,
+						refreshToken: mockRefreshToken,
+						expiresOn: ''
+					});
+				});
+				delete (mockAccessToken as any).tokenType;
+
+				azureAuthCodeGrant.setup(x => x.refreshToken(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
+					return Promise.resolve({
+						accessToken: mockAccessToken
+					} as OAuthTokenResponse);
+				});
+
+				const securityToken: Token = await azureAuthCodeGrant.object.getAccountSecurityToken(mockAccount, mockTenant.id, AzureResource.MicrosoftResourceManagement);
+				should(securityToken.tokenType).be.equal('Bearer', 'tokenType should be bearer on a successful getSecurityToken from cache')
+			});
+		});

-		await baseAuth.clearCredentials(accountKey);
-		result = await baseAuth.getCachedToken(accountKey, resourceId, tenantId);
-		should(result).be.undefined();
 	});

-	it('Token set with resource ID and get without tenant and resource id', async function (): Promise<void> {
-		await baseAuth.setCachedToken(accountKey, accessToken, refreshToken, resourceId, tenantId);
-		const result = await baseAuth.getCachedToken(accountKey);
+	describe('getToken', function () {

-		should(JSON.stringify(result)).be.undefined();
-		should(JSON.stringify(result)).be.undefined();
+		it('calls handle interaction required', async function () {
+			azureAuthCodeGrant.setup(x => x.makePostRequest(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
+				return Promise.resolve({
+					data: {
+						error: 'interaction_required'
+					}
+				} as AxiosResponse<any>);
+			});
+
+			azureAuthCodeGrant.setup(x => x.handleInteractionRequired(mockTenant, provider.settings.microsoftResource)).returns(() => {
+				return Promise.resolve({
+					accessToken: mockAccessToken
+				} as OAuthTokenResponse);
+			});
+
+
+			const result = await azureAuthCodeGrant.object.getToken(mockTenant, provider.settings.microsoftResource, {} as TokenPostData);
+
+			azureAuthCodeGrant.verify(x => x.handleInteractionRequired(mockTenant, provider.settings.microsoftResource), TypeMoq.Times.once());
+
+			should(result.accessToken).be.deepEqual(mockAccessToken);
+		});
+
+		it('unknown error should throw error', async function () {
+			azureAuthCodeGrant.setup(x => x.makePostRequest(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
+				return Promise.resolve({
+					data: {
+						error: 'unknown error'
+					}
+				} as AxiosResponse<any>);
+			});
+
+			await azureAuthCodeGrant.object.getToken(mockTenant, provider.settings.microsoftResource, {} as TokenPostData).should.be.rejected();
+		});
+
+		it('calls getTokenHelper', async function () {
+			azureAuthCodeGrant.setup(x => x.makePostRequest(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
+				return Promise.resolve({
+					data: {
+						access_token: mockAccessToken.token,
+						refresh_token: mockRefreshToken.token,
+						expires_on: `0`
+					}
+				} as AxiosResponse<any>);
+			});
+
+			azureAuthCodeGrant.setup(x => x.getTokenHelper(mockTenant, provider.settings.microsoftResource, TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
+				return Promise.resolve({
+					accessToken: mockAccessToken
+				} as OAuthTokenResponse);
+			});
+
+
+			const result = await azureAuthCodeGrant.object.getToken(mockTenant, provider.settings.microsoftResource, {} as TokenPostData);
+
+			azureAuthCodeGrant.verify(x => x.getTokenHelper(mockTenant, provider.settings.microsoftResource, TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.once());
+
+			should(result.accessToken).be.deepEqual(mockAccessToken);
+		});
 	});

-	it('Create an account object', async function (): Promise<void> {
-		const tokenClaims = {
-			idp: 'live.com',
-			name: 'TestAccount',
-		} as TokenClaims;
-
-		const account = baseAuth.createAccount(tokenClaims, 'someKey', undefined);
-
-		should(account.properties.azureAuthType).be.equal(AzureAuthType.AuthCodeGrant);
-		should(account.key.accountId).be.equal('someKey');
-		should(account.properties.isMsAccount).be.equal(true);
-	});
-	it('Should handle ignored tenants', async function (): Promise<void> {
-		// Don't sit on the await openConsentDialog if test is failing
-		this.timeout(3000);
-
-		const configuration = vscode.workspace.getConfiguration('azure.tenant.config');
-		const values = [tenantId];
-
-		await configuration.update('filter', values, vscode.ConfigurationTarget.Global);
-		const x = await baseAuth.openConsentDialog(tenant, resourceId);
-
-		should(x).be.false();
-	});
 });
--- a/extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts
+++ b/extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts
@@ -41,12 +41,15 @@ const mockAccount: AzureAccount = {
 	isStale: false
 };

+const mockTenantId: string = 'mock_tenant';
+
+
 const mockSubscription: azureResource.AzureResourceSubscription = {
 	id: 'mock_subscription',
-	name: 'mock subscription'
+	name: 'mock subscription',
+	tenant: mockTenantId
 };

-const mockTenantId: string = 'mock_tenant';

 const mockResourceRootNode: azureResource.IAzureResourceNode = {
 	account: mockAccount,
@@ -61,8 +64,7 @@ const mockResourceRootNode: azureResource.IAzureResourceNode = {
 	}
 };

-const mockTokens: { [key: string]: any } = {};
-mockTokens[mockTenantId] = {
+const mockToken = {
 	token: 'mock_token',
 	tokenType: 'Bearer'
 };
@@ -106,7 +108,7 @@ describe('AzureResourceDatabaseTreeDataProvider.getChildren', function (): void
 		mockDatabaseService = TypeMoq.Mock.ofType<IAzureResourceService<azureResource.AzureResourceDatabase>>();
 		mockExtensionContext = TypeMoq.Mock.ofType<vscode.ExtensionContext>();

-		sinon.stub(azdata.accounts, 'getSecurityToken').returns(Promise.resolve(mockTokens));
+		sinon.stub(azdata.accounts, 'getAccountSecurityToken').returns(Promise.resolve(mockToken));
 		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());
 	});
--- a/extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts
+++ b/extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts
@@ -41,12 +41,14 @@ const mockAccount: AzureAccount = {
 	isStale: false
 };

+const mockTenantId: string = 'mock_tenant';
+
 const mockSubscription: azureResource.AzureResourceSubscription = {
 	id: 'mock_subscription',
-	name: 'mock subscription'
+	name: 'mock subscription',
+	tenant: mockTenantId
 };

-const mockTenantId: string = 'mock_tenant';

 const mockResourceRootNode: azureResource.IAzureResourceNode = {
 	account: mockAccount,
@@ -61,8 +63,7 @@ const mockResourceRootNode: azureResource.IAzureResourceNode = {
 	}
 };

-const mockTokens: { [key: string]: any } = {};
-mockTokens[mockTenantId] = {
+const mockToken = {
 	token: 'mock_token',
 	tokenType: 'Bearer'
 };
@@ -106,7 +107,7 @@ describe('AzureResourceDatabaseServerTreeDataProvider.getChildren', function ():
 		mockDatabaseServerService = TypeMoq.Mock.ofType<IAzureResourceService<azureResource.AzureResourceDatabaseServer>>();
 		mockExtensionContext = TypeMoq.Mock.ofType<vscode.ExtensionContext>();

-		sinon.stub(azdata.accounts, 'getSecurityToken').returns(Promise.resolve(mockTokens));
+		sinon.stub(azdata.accounts, 'getAccountSecurityToken').returns(Promise.resolve(mockToken));
 		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());
 	});
--- a/extensions/azurecore/src/test/azureResource/resourceService.test.ts
+++ b/extensions/azurecore/src/test/azureResource/resourceService.test.ts
@@ -33,13 +33,14 @@ const mockAccount: AzureAccount = {
 	isStale: false
 };

+const mockTenantId: string = 'mock_tenant';
+
 const mockSubscription: azureResource.AzureResourceSubscription = {
 	id: 'mock_subscription',
-	name: 'mock subscription'
+	name: 'mock subscription',
+	tenant: mockTenantId
 };

-const mockTenantId: string = 'mock_tenant';
-
 let mockResourceTreeDataProvider1: TypeMoq.IMock<azureResource.IAzureResourceTreeDataProvider>;
 let mockResourceProvider1: TypeMoq.IMock<azureResource.IAzureResourceProvider>;

--- a/extensions/azurecore/src/test/azureResource/resourceTreeNode.test.ts
+++ b/extensions/azurecore/src/test/azureResource/resourceTreeNode.test.ts
@@ -36,12 +36,14 @@ const mockAccount: AzureAccount = {
 	isStale: false
 };

+const mockTenantId: string = 'mock_tenant';
+
 const mockSubscription: azureResource.AzureResourceSubscription = {
 	id: 'mock_subscription',
-	name: 'mock subscription'
+	name: 'mock subscription',
+	tenant: mockTenantId
 };

-const mockTenantId: string = 'mock_tenant';

 const mockResourceProviderId: string = 'mock_resource_provider';

--- a/extensions/azurecore/src/test/azureResource/tree/accountTreeNode.test.ts
+++ b/extensions/azurecore/src/test/azureResource/tree/accountTreeNode.test.ts
@@ -17,7 +17,6 @@ import {
 	IAzureResourceCacheService,
 	IAzureResourceSubscriptionService,
 	IAzureResourceSubscriptionFilterService,
-	IAzureResourceTenantService
 } from '../../../azureResource/interfaces';
 import { IAzureResourceTreeChangeHandler } from '../../../azureResource/tree/treeChangeHandler';
 import { AzureResourceAccountTreeNode } from '../../../azureResource/tree/accountTreeNode';
@@ -31,7 +30,6 @@ let mockExtensionContext: TypeMoq.IMock<vscode.ExtensionContext>;
 let mockCacheService: TypeMoq.IMock<IAzureResourceCacheService>;
 let mockSubscriptionService: TypeMoq.IMock<IAzureResourceSubscriptionService>;
 let mockSubscriptionFilterService: TypeMoq.IMock<IAzureResourceSubscriptionFilterService>;
-let mockTenantService: TypeMoq.IMock<IAzureResourceTenantService>;
 let mockAppContext: AppContext;
 let getSecurityTokenStub: sinon.SinonStub;
 let mockTreeChangeHandler: TypeMoq.IMock<IAzureResourceTreeChangeHandler>;
@@ -63,28 +61,27 @@ const mockAccount: azdata.Account = {

 const mockSubscription1: azureResource.AzureResourceSubscription = {
 	id: 'mock_subscription_1',
-	name: 'mock subscription 1'
+	name: 'mock subscription 1',
+	tenant: mockTenantId
 };

 const mockSubscription2: azureResource.AzureResourceSubscription = {
 	id: 'mock_subscription_2',
-	name: 'mock subscription 2'
+	name: 'mock subscription 2',
+	tenant: mockTenantId
 };

 const mockSubscriptions = [mockSubscription1, mockSubscription2];

 const mockFilteredSubscriptions = [mockSubscription1];

-const mockTokens: { [key: string]: any } = {};
+const mockToken = {
+	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);
+const mockCredential = new TokenCredentials(mockToken.token, mockToken.tokenType);

 let mockSubscriptionCache: azureResource.AzureResourceSubscription[] = [];

@@ -94,7 +91,6 @@ describe('AzureResourceAccountTreeNode.info', function (): void {
 		mockCacheService = TypeMoq.Mock.ofType<IAzureResourceCacheService>();
 		mockSubscriptionService = TypeMoq.Mock.ofType<IAzureResourceSubscriptionService>();
 		mockSubscriptionFilterService = TypeMoq.Mock.ofType<IAzureResourceSubscriptionFilterService>();
-		mockTenantService = TypeMoq.Mock.ofType<IAzureResourceTenantService>();

 		mockTreeChangeHandler = TypeMoq.Mock.ofType<IAzureResourceTreeChangeHandler>();

@@ -104,13 +100,11 @@ describe('AzureResourceAccountTreeNode.info', function (): void {
 		mockAppContext.registerService<IAzureResourceCacheService>(AzureResourceServiceNames.cacheService, mockCacheService.object);
 		mockAppContext.registerService<IAzureResourceSubscriptionService>(AzureResourceServiceNames.subscriptionService, mockSubscriptionService.object);
 		mockAppContext.registerService<IAzureResourceSubscriptionFilterService>(AzureResourceServiceNames.subscriptionFilterService, mockSubscriptionFilterService.object);
-		mockAppContext.registerService<IAzureResourceTenantService>(AzureResourceServiceNames.tenantService, mockTenantService.object);

-		getSecurityTokenStub = sinon.stub(azdata.accounts, 'getSecurityToken').returns(Promise.resolve(mockTokens));
+		getSecurityTokenStub = sinon.stub(azdata.accounts, 'getAccountSecurityToken').returns(Promise.resolve(mockToken));
 		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(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
 	});

 	afterEach(function (): void {
@@ -138,7 +132,7 @@ describe('AzureResourceAccountTreeNode.info', function (): void {
 	});

 	it('Should be correct when there are subscriptions listed.', async function (): Promise<void> {
-		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential)).returns(() => Promise.resolve(mockSubscriptions));
+		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId)).returns(() => Promise.resolve(mockSubscriptions));
 		mockSubscriptionFilterService.setup((o) => o.getSelectedSubscriptions(mockAccount)).returns(() => Promise.resolve(undefined));

 		const accountTreeNodeLabel = `${mockAccount.displayInfo.displayName} (${mockSubscriptions.length} / ${mockSubscriptions.length} subscriptions)`;
@@ -158,7 +152,7 @@ describe('AzureResourceAccountTreeNode.info', function (): void {
 	});

 	it('Should be correct when there are subscriptions filtered.', async function (): Promise<void> {
-		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential)).returns(() => Promise.resolve(mockSubscriptions));
+		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId)).returns(() => Promise.resolve(mockSubscriptions));
 		mockSubscriptionFilterService.setup((o) => o.getSelectedSubscriptions(mockAccount)).returns(() => Promise.resolve(mockFilteredSubscriptions));

 		const accountTreeNodeLabel = `${mockAccount.displayInfo.displayName} (${mockFilteredSubscriptions.length} / ${mockSubscriptions.length} subscriptions)`;
@@ -184,7 +178,6 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 		mockCacheService = TypeMoq.Mock.ofType<IAzureResourceCacheService>();
 		mockSubscriptionService = TypeMoq.Mock.ofType<IAzureResourceSubscriptionService>();
 		mockSubscriptionFilterService = TypeMoq.Mock.ofType<IAzureResourceSubscriptionFilterService>();
-		mockTenantService = TypeMoq.Mock.ofType<IAzureResourceTenantService>();

 		mockTreeChangeHandler = TypeMoq.Mock.ofType<IAzureResourceTreeChangeHandler>();

@@ -194,13 +187,11 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 		mockAppContext.registerService<IAzureResourceCacheService>(AzureResourceServiceNames.cacheService, mockCacheService.object);
 		mockAppContext.registerService<IAzureResourceSubscriptionService>(AzureResourceServiceNames.subscriptionService, mockSubscriptionService.object);
 		mockAppContext.registerService<IAzureResourceSubscriptionFilterService>(AzureResourceServiceNames.subscriptionFilterService, mockSubscriptionFilterService.object);
-		mockAppContext.registerService<IAzureResourceTenantService>(AzureResourceServiceNames.tenantService, mockTenantService.object);

-		sinon.stub(azdata.accounts, 'getSecurityToken').returns(Promise.resolve(mockTokens));
+		sinon.stub(azdata.accounts, 'getAccountSecurityToken').returns(Promise.resolve(mockToken));
 		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(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
 	});

 	afterEach(function (): void {
@@ -208,14 +199,14 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 	});

 	it('Should load subscriptions from scratch and update cache when it is clearing cache.', async function (): Promise<void> {
-		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential)).returns(() => Promise.resolve(mockSubscriptions));
+		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId)).returns(() => Promise.resolve(mockSubscriptions));
 		mockSubscriptionFilterService.setup((o) => o.getSelectedSubscriptions(mockAccount)).returns(() => Promise.resolve([]));

 		const accountTreeNode = new AzureResourceAccountTreeNode(mockAccount, mockAppContext, mockTreeChangeHandler.object);

 		const children = await accountTreeNode.getChildren();

-		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential), TypeMoq.Times.once());
+		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId), 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());
 		mockSubscriptionFilterService.verify((o) => o.getSelectedSubscriptions(mockAccount), TypeMoq.Times.once());
@@ -241,7 +232,7 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 	});

 	it('Should load subscriptions from cache when it is not clearing cache.', async function (): Promise<void> {
-		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential)).returns(() => Promise.resolve(mockSubscriptions));
+		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId)).returns(() => Promise.resolve(mockSubscriptions));
 		mockSubscriptionFilterService.setup((o) => o.getSelectedSubscriptions(mockAccount)).returns(() => Promise.resolve(undefined));

 		const accountTreeNode = new AzureResourceAccountTreeNode(mockAccount, mockAppContext, mockTreeChangeHandler.object);
@@ -250,7 +241,7 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 		const children = await accountTreeNode.getChildren();


-		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential), TypeMoq.Times.once());
+		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId), 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());

@@ -262,7 +253,7 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 	});

 	it('Should handle when there is no subscriptions.', async function (): Promise<void> {
-		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential)).returns(() => Promise.resolve(undefined));
+		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId)).returns(() => Promise.resolve(undefined));

 		const accountTreeNode = new AzureResourceAccountTreeNode(mockAccount, mockAppContext, mockTreeChangeHandler.object);

@@ -278,7 +269,7 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 	});

 	it('Should honor subscription filtering.', async function (): Promise<void> {
-		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential)).returns(() => Promise.resolve(mockSubscriptions));
+		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId)).returns(() => Promise.resolve(mockSubscriptions));
 		mockSubscriptionFilterService.setup((o) => o.getSelectedSubscriptions(mockAccount)).returns(() => Promise.resolve(mockFilteredSubscriptions));

 		const accountTreeNode = new AzureResourceAccountTreeNode(mockAccount, mockAppContext, mockTreeChangeHandler.object);
@@ -296,7 +287,7 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {
 	});

 	it('Should handle errors.', async function (): Promise<void> {
-		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential)).returns(() => Promise.resolve(mockSubscriptions));
+		mockSubscriptionService.setup((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId)).returns(() => Promise.resolve(mockSubscriptions));

 		const mockError = 'Test error';
 		mockSubscriptionFilterService.setup((o) => o.getSelectedSubscriptions(mockAccount)).returns(() => { throw new Error(mockError); });
@@ -305,8 +296,8 @@ describe('AzureResourceAccountTreeNode.getChildren', function (): void {

 		const children = await accountTreeNode.getChildren();

-		should(getSecurityTokenStub.calledOnce).be.true('getSecurityToken should have been called exactly once');
-		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential), TypeMoq.Times.once());
+		should(getSecurityTokenStub.calledTwice).be.true('getSecurityToken should have been called exactly twice - once per subscription');
+		mockSubscriptionService.verify((o) => o.getSubscriptions(mockAccount, mockCredential, mockTenantId), TypeMoq.Times.once());
 		mockSubscriptionFilterService.verify((o) => o.getSelectedSubscriptions(mockAccount), TypeMoq.Times.once());
 		mockCacheService.verify((o) => o.get(TypeMoq.It.isAnyString()), TypeMoq.Times.never());
 		mockCacheService.verify((o) => o.update(TypeMoq.It.isAnyString(), TypeMoq.It.isAny()), TypeMoq.Times.once());
@@ -325,7 +316,6 @@ describe('AzureResourceAccountTreeNode.clearCache', function (): void {
 		mockCacheService = TypeMoq.Mock.ofType<IAzureResourceCacheService>();
 		mockSubscriptionService = TypeMoq.Mock.ofType<IAzureResourceSubscriptionService>();
 		mockSubscriptionFilterService = TypeMoq.Mock.ofType<IAzureResourceSubscriptionFilterService>();
-		mockTenantService = TypeMoq.Mock.ofType<IAzureResourceTenantService>();

 		mockTreeChangeHandler = TypeMoq.Mock.ofType<IAzureResourceTreeChangeHandler>();

@@ -335,13 +325,11 @@ describe('AzureResourceAccountTreeNode.clearCache', function (): void {
 		mockAppContext.registerService<IAzureResourceCacheService>(AzureResourceServiceNames.cacheService, mockCacheService.object);
 		mockAppContext.registerService<IAzureResourceSubscriptionService>(AzureResourceServiceNames.subscriptionService, mockSubscriptionService.object);
 		mockAppContext.registerService<IAzureResourceSubscriptionFilterService>(AzureResourceServiceNames.subscriptionFilterService, mockSubscriptionFilterService.object);
-		mockAppContext.registerService<IAzureResourceTenantService>(AzureResourceServiceNames.tenantService, mockTenantService.object);

-		sinon.stub(azdata.accounts, 'getSecurityToken').returns(Promise.resolve(mockTokens));
+		sinon.stub(azdata.accounts, 'getAccountSecurityToken').returns(Promise.resolve(mockToken));
 		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(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(mockTenantId));
 	});

 	afterEach(function (): void {
--- a/extensions/azurecore/src/test/azureResource/tree/subscriptionTreeNode.test.ts
+++ b/extensions/azurecore/src/test/azureResource/tree/subscriptionTreeNode.test.ts
@@ -43,13 +43,14 @@ const mockAccount: azdata.Account = {
 	isStale: false
 };

+const mockTenantId: string = 'mock_tenant';
+
 const mockSubscription: azureResource.AzureResourceSubscription = {
 	id: 'mock_subscription',
-	name: 'mock subscription'
+	name: 'mock subscription',
+	tenant: mockTenantId
 };

-const mockTenantId: string = 'mock_tenant';
-
 let mockResourceTreeDataProvider1: TypeMoq.IMock<azureResource.IAzureResourceTreeDataProvider>;
 let mockResourceProvider1: TypeMoq.IMock<azureResource.IAzureResourceProvider>;