From 8f75138ea3e131449007f6f672a8aa85048058ca Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 8 May 2025 14:05:41 -0400 Subject: [PATCH 1/7] test(amazonq): add unit tests for authUtil --- .../unit/codewhisperer/util/authUtil.test.ts | 610 ++++++------------ 1 file changed, 193 insertions(+), 417 deletions(-) diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index e4f73c4df05..8ed2188618b 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts @@ -3,420 +3,196 @@ * SPDX-License-Identifier: Apache-2.0 */ -// import assert from 'assert' -// import { -// AuthStates, -// AuthUtil, -// amazonQScopes, -// codeWhispererChatScopes, -// codeWhispererCoreScopes, -// } from 'aws-core-vscode/codewhisperer' -// import { -// assertTelemetry, -// getTestWindow, -// SeverityLevel, -// createBuilderIdProfile, -// createSsoProfile, -// createTestAuth, -// captureEventNTimes, -// } from 'aws-core-vscode/test' -// import { Auth, Connection, isAnySsoConnection, isBuilderIdConnection } from 'aws-core-vscode/auth' -// import { globals, vscodeComponent } from 'aws-core-vscode/shared' - -// const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start' - -// describe('AuthUtil', async function () { -// let auth: ReturnType -// let authUtil: AuthUtil - -// beforeEach(async function () { -// auth = createTestAuth(globals.globalState) -// authUtil = new AuthUtil(auth) -// }) - -// afterEach(async function () { -// await auth.logout() -// }) - -// it('if there is no valid AwsBuilderID conn, it will create one and use it', async function () { -// getTestWindow().onDidShowQuickPick(async (picker) => { -// await picker.untilReady() -// picker.acceptItem(picker.items[1]) -// }) - -// await authUtil.connectToAwsBuilderId() -// const conn = authUtil.conn -// assert.strictEqual(conn?.type, 'sso') -// assert.strictEqual(conn.label, 'AWS Builder ID') -// assert.deepStrictEqual(conn.scopes, amazonQScopes) -// }) - -// it('if there IS an existing AwsBuilderID conn, it will upgrade the scopes and use it', async function () { -// const existingBuilderId = await auth.createConnection( -// createBuilderIdProfile({ scopes: codeWhispererCoreScopes }) -// ) -// getTestWindow().onDidShowQuickPick(async (picker) => { -// await picker.untilReady() -// picker.acceptItem(picker.items[1]) -// }) - -// await authUtil.connectToAwsBuilderId() - -// const conn = authUtil.conn -// assert.strictEqual(conn?.type, 'sso') -// assert.strictEqual(conn.id, existingBuilderId.id) -// assert.deepStrictEqual(conn.scopes, amazonQScopes) -// }) - -// it('if there is no valid enterprise SSO conn, will create and use one', async function () { -// getTestWindow().onDidShowQuickPick(async (picker) => { -// await picker.untilReady() -// picker.acceptItem(picker.items[1]) -// }) - -// await authUtil.connectToEnterpriseSso(enterpriseSsoStartUrl, 'us-east-1') -// const conn = authUtil.conn -// assert.strictEqual(conn?.type, 'sso') -// assert.strictEqual(conn.label, 'IAM Identity Center (enterprise)') -// }) - -// it('should add scopes + connect to existing IAM Identity Center connection', async function () { -// getTestWindow().onDidShowMessage(async (message) => { -// assert.ok(message.modal) -// message.selectItem('Proceed') -// }) -// const randomScope = 'my:random:scope' -// const ssoConn = await auth.createInvalidSsoConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: [randomScope] }) -// ) - -// // Method under test -// await authUtil.connectToEnterpriseSso(ssoConn.startUrl, 'us-east-1') - -// const cwConn = authUtil.conn -// assert.strictEqual(cwConn?.type, 'sso') -// assert.strictEqual(cwConn.label, 'IAM Identity Center (enterprise)') -// assert.deepStrictEqual(cwConn.scopes, [randomScope, ...amazonQScopes]) -// }) - -// it('reauthenticates an existing BUT invalid Amazon Q IAM Identity Center connection', async function () { -// const ssoConn = await auth.createInvalidSsoConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) -// ) -// await auth.refreshConnectionState(ssoConn) -// assert.strictEqual(auth.getConnectionState(ssoConn), 'invalid') - -// // Method under test -// await authUtil.connectToEnterpriseSso(ssoConn.startUrl, 'us-east-1') - -// const cwConn = authUtil.conn -// assert.strictEqual(cwConn?.type, 'sso') -// assert.strictEqual(cwConn.id, ssoConn.id) -// assert.deepStrictEqual(cwConn.scopes, amazonQScopes) -// assert.strictEqual(auth.getConnectionState(cwConn), 'valid') -// }) - -// it('should show reauthenticate prompt', async function () { -// getTestWindow().onDidShowMessage((m) => { -// if (m.severity === SeverityLevel.Information) { -// m.close() -// } -// }) - -// await authUtil.showReauthenticatePrompt() - -// const warningMessage = getTestWindow().shownMessages.filter((m) => m.severity === SeverityLevel.Information) -// assert.strictEqual(warningMessage.length, 1) -// assert.strictEqual(warningMessage[0].message, `Your Amazon Q connection has expired. Please re-authenticate.`) -// warningMessage[0].close() -// assertTelemetry('toolkit_showNotification', { -// id: 'codeWhispererConnectionExpired', -// result: 'Succeeded', -// source: vscodeComponent, -// }) -// assertTelemetry('toolkit_invokeAction', { -// id: 'codeWhispererConnectionExpired', -// action: 'dismiss', -// result: 'Succeeded', -// source: vscodeComponent, -// }) -// }) - -// it('reauthenticate prompt reauthenticates invalid connection', async function () { -// const conn = await auth.createInvalidSsoConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererChatScopes }) -// ) -// await auth.useConnection(conn) -// getTestWindow().onDidShowMessage((m) => { -// m.selectItem('Re-authenticate') -// }) - -// assert.strictEqual(auth.getConnectionState(conn), 'invalid') - -// await authUtil.showReauthenticatePrompt() - -// assert.strictEqual(authUtil.conn?.type, 'sso') -// assert.strictEqual(auth.getConnectionState(conn), 'valid') -// assertTelemetry('toolkit_showNotification', { -// id: 'codeWhispererConnectionExpired', -// result: 'Succeeded', -// source: vscodeComponent, -// }) -// assertTelemetry('toolkit_invokeAction', { -// id: 'codeWhispererConnectionExpired', -// action: 'connect', -// result: 'Succeeded', -// source: vscodeComponent, -// }) -// }) - -// it('reauthenticates Builder ID connection that already has all scopes', async function () { -// const conn = await auth.createInvalidSsoConnection(createBuilderIdProfile({ scopes: amazonQScopes })) -// await auth.useConnection(conn) - -// // method under test -// await authUtil.reauthenticate() - -// assert.strictEqual(authUtil.conn?.type, 'sso') -// assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) -// assert.strictEqual(auth.getConnectionState(conn), 'valid') -// }) - -// it('reauthenticates IdC connection that already has all scopes', async function () { -// const conn = await auth.createInvalidSsoConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererCoreScopes }) -// ) -// await auth.useConnection(conn) - -// // method under test -// await authUtil.reauthenticate() - -// assert.strictEqual(authUtil.conn?.type, 'sso') -// assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) -// assert.strictEqual(auth.getConnectionState(conn), 'valid') -// }) - -// it('reauthenticate adds missing Builder ID scopes', async function () { -// const conn = await auth.createInvalidSsoConnection(createBuilderIdProfile({ scopes: codeWhispererCoreScopes })) -// await auth.useConnection(conn) - -// // method under test -// await authUtil.reauthenticate() - -// assert.strictEqual(authUtil.conn?.type, 'sso') -// assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) -// assert.strictEqual(auth.getConnectionState(conn), 'valid') -// }) - -// it('reauthenticate adds missing Amazon Q IdC scopes', async function () { -// const conn = await auth.createInvalidSsoConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererCoreScopes }) -// ) -// await auth.useConnection(conn) - -// // method under test -// await authUtil.reauthenticate() - -// assert.strictEqual(authUtil.conn?.type, 'sso') -// assert.deepStrictEqual(authUtil.conn?.scopes, amazonQScopes) -// assert.strictEqual(auth.getConnectionState(conn), 'valid') -// }) - -// it('CodeWhisperer uses fallback connection when switching to an unsupported connection', async function () { -// const supportedConn = await auth.createConnection(createBuilderIdProfile({ scopes: codeWhispererChatScopes })) -// const unsupportedConn = await auth.createConnection(createSsoProfile()) - -// await auth.useConnection(supportedConn) -// assert.ok(authUtil.isConnected()) -// assert.strictEqual(auth.activeConnection?.id, authUtil.conn?.id) - -// // Switch to unsupported connection -// const cwAuthUpdatedConnection = captureEventNTimes(authUtil.secondaryAuth.onDidChangeActiveConnection, 2) -// await auth.useConnection(unsupportedConn) -// // - This is triggered when the main Auth connection is switched -// // - This is triggered by registerAuthListener() when it saves the previous active connection as a fallback. -// await cwAuthUpdatedConnection - -// // TODO in a refactor see if we can simplify multiple multiple triggers on the same event. -// assert.ok(authUtil.isConnected()) -// assert.ok(authUtil.isUsingSavedConnection) -// assert.notStrictEqual(auth.activeConnection?.id, authUtil.conn?.id) -// assert.strictEqual(authUtil.conn?.type, 'sso') -// assert.deepStrictEqual(authUtil.conn?.scopes, codeWhispererChatScopes) -// }) - -// it('does not prompt to sign out of duplicate builder ID connections', async function () { -// await authUtil.connectToAwsBuilderId() -// await authUtil.connectToAwsBuilderId() -// assert.ok(authUtil.isConnected()) - -// const ssoConnectionIds = new Set(auth.activeConnectionEvents.emits.filter(isAnySsoConnection).map((c) => c.id)) -// assert.strictEqual(ssoConnectionIds.size, 1, 'Expected exactly 1 unique SSO connection id') -// assert.strictEqual((await auth.listConnections()).filter(isAnySsoConnection).length, 1) -// }) - -// it('automatically upgrades connections if they do not have the required scopes', async function () { -// const upgradeableConn = await auth.createConnection(createBuilderIdProfile()) -// await auth.useConnection(upgradeableConn) -// assert.strictEqual(authUtil.isConnected(), false) - -// await authUtil.connectToAwsBuilderId() -// assert.ok(authUtil.isConnected()) -// assert.ok(authUtil.isConnectionValid()) -// assert.ok(isBuilderIdConnection(authUtil.conn)) -// assert.strictEqual(authUtil.conn?.id, upgradeableConn.id) -// assert.strictEqual(authUtil.conn.startUrl, upgradeableConn.startUrl) -// assert.strictEqual(authUtil.conn.ssoRegion, upgradeableConn.ssoRegion) -// assert.deepStrictEqual(authUtil.conn.scopes, amazonQScopes) -// assert.strictEqual((await auth.listConnections()).filter(isAnySsoConnection).length, 1) -// }) - -// it('test reformatStartUrl should remove trailing slash and hash', function () { -// const expected = 'https://view.awsapps.com/start' -// assert.strictEqual(authUtil.reformatStartUrl(expected + '/'), expected) -// assert.strictEqual(authUtil.reformatStartUrl(undefined), undefined) -// assert.strictEqual(authUtil.reformatStartUrl(expected + '/#'), expected) -// assert.strictEqual(authUtil.reformatStartUrl(expected + '#/'), expected) -// assert.strictEqual(authUtil.reformatStartUrl(expected + '/#/'), expected) -// assert.strictEqual(authUtil.reformatStartUrl(expected + '####'), expected) -// }) - -// it(`clearExtraConnections()`, async function () { -// const conn1 = await auth.createConnection(createBuilderIdProfile()) -// const conn2 = await auth.createConnection(createSsoProfile({ startUrl: enterpriseSsoStartUrl })) -// const conn3 = await auth.createConnection(createSsoProfile({ startUrl: enterpriseSsoStartUrl + 1 })) -// // validate listConnections shows all connections -// assert.deepStrictEqual( -// (await authUtil.auth.listConnections()).map((conn) => conn.id).sort((a, b) => a.localeCompare(b)), -// [conn1, conn2, conn3].map((conn) => conn.id).sort((a, b) => a.localeCompare(b)) -// ) -// await authUtil.secondaryAuth.useNewConnection(conn3) - -// await authUtil.clearExtraConnections() // method under test - -// // Only the conn that AuthUtil is using is remaining -// assert.deepStrictEqual( -// (await authUtil.auth.listConnections()).map((conn) => conn.id), -// [conn3.id] -// ) -// }) -// }) - -// describe('getChatAuthState()', function () { -// let auth: ReturnType -// let authUtil: AuthUtil -// let laterDate: Date - -// beforeEach(async function () { -// auth = createTestAuth(globals.globalState) -// authUtil = new AuthUtil(auth) - -// laterDate = new Date(Date.now() + 10_000_000) -// }) - -// afterEach(async function () { -// await auth.logout() -// }) - -// it('indicates nothing connected when no auth connection exists', async function () { -// const result = await authUtil.getChatAuthState() -// assert.deepStrictEqual(result, { -// codewhispererChat: AuthStates.disconnected, -// codewhispererCore: AuthStates.disconnected, -// amazonQ: AuthStates.disconnected, -// }) -// }) - -// /** Affects {@link Auth.refreshConnectionState} */ -// function createToken(conn: Connection) { -// auth.getTestTokenProvider(conn).getToken.resolves({ accessToken: 'myAccessToken', expiresAt: laterDate }) -// } - -// describe('Builder ID', function () { -// it('indicates only CodeWhisperer core is connected when only CW core scopes are set', async function () { -// const conn = await auth.createConnection(createBuilderIdProfile({ scopes: codeWhispererCoreScopes })) -// createToken(conn) -// await auth.useConnection(conn) - -// const result = await authUtil.getChatAuthState() -// assert.deepStrictEqual(result, { -// codewhispererCore: AuthStates.connected, -// codewhispererChat: AuthStates.expired, -// amazonQ: AuthStates.expired, -// }) -// }) - -// it('indicates all SUPPORTED features connected when all scopes are set', async function () { -// const conn = await auth.createConnection(createBuilderIdProfile({ scopes: amazonQScopes })) -// createToken(conn) -// await auth.useConnection(conn) - -// const result = await authUtil.getChatAuthState() -// assert.deepStrictEqual(result, { -// codewhispererCore: AuthStates.connected, -// codewhispererChat: AuthStates.connected, -// amazonQ: AuthStates.connected, -// }) -// }) - -// it('indicates all SUPPORTED features expired when connection is invalid', async function () { -// const conn = await auth.createInvalidSsoConnection( -// createBuilderIdProfile({ scopes: codeWhispererChatScopes }) -// ) -// await auth.useConnection(conn) - -// const result = await authUtil.getChatAuthState() -// assert.deepStrictEqual(result, { -// codewhispererCore: AuthStates.expired, -// codewhispererChat: AuthStates.expired, -// amazonQ: AuthStates.expired, -// }) -// }) -// }) - -// describe('Identity Center', function () { -// it('indicates only CW core is connected when only CW core scopes are set', async function () { -// const conn = await auth.createConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: codeWhispererCoreScopes }) -// ) -// createToken(conn) -// await auth.useConnection(conn) - -// const result = await authUtil.getChatAuthState() -// assert.deepStrictEqual(result, { -// codewhispererCore: AuthStates.pendingProfileSelection, -// codewhispererChat: AuthStates.expired, -// amazonQ: AuthStates.expired, -// }) -// }) - -// it('indicates all features connected when all scopes are set', async function () { -// const conn = await auth.createConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) -// ) -// createToken(conn) -// await auth.useConnection(conn) - -// const result = await authUtil.getChatAuthState() -// assert.deepStrictEqual(result, { -// codewhispererCore: AuthStates.pendingProfileSelection, -// codewhispererChat: AuthStates.pendingProfileSelection, -// amazonQ: AuthStates.pendingProfileSelection, -// }) -// }) - -// it('indicates all features expired when connection is invalid', async function () { -// const conn = await auth.createInvalidSsoConnection( -// createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) -// ) -// await auth.useConnection(conn) - -// const result = await authUtil.getChatAuthState() -// assert.deepStrictEqual(result, { -// codewhispererCore: AuthStates.expired, -// codewhispererChat: AuthStates.expired, -// amazonQ: AuthStates.expired, -// }) -// }) -// }) -// }) +import assert from 'assert' +import * as sinon from 'sinon' +import { AuthUtil } from 'aws-core-vscode/codewhisperer' +import { createTestAuthUtil } from 'aws-core-vscode/test' +import { constants } from 'aws-core-vscode/auth' +import { auth2 } from 'aws-core-vscode/auth' + +describe('AuthUtil', async function () { + await createTestAuthUtil() + + const auth = AuthUtil.instance + const originalSession = (auth as any).session + + beforeEach(async function () { + ;(auth as any).session = originalSession + }) + + afterEach(async function () { + sinon.restore() + }) + + describe('Auth state', function () { + it('login with BuilderId', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + assert.ok(auth.isConnected()) + assert.ok(auth.isBuilderIdConnection()) + }) + + it('login with IDC', async function () { + await auth.login('https://example.awsapps.com/start', 'us-east-1') + assert.ok(auth.isConnected()) + assert.ok(auth.isIdcConnection()) + }) + + it('identifies internal users', async function () { + await auth.login(constants.internalStartUrl, 'us-east-1') + assert.ok(auth.isInternalAmazonUser()) + }) + + it('identifies SSO session', function () { + ;(auth as any).session = { loginType: auth2.LoginTypes.SSO } + assert.strictEqual(auth.isSsoSession(), true) + }) + + it('identifies non-SSO session', function () { + ;(auth as any).session = { loginType: auth2.LoginTypes.IAM } + assert.strictEqual(auth.isSsoSession(), false) + }) + }) + + describe('Token management', function () { + it('can get token when connected with SSO', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + const token = await auth.getToken() + assert.ok(token) + }) + + it('throws when getting token without SSO connection', async function () { + sinon.stub(AuthUtil.instance, 'isSsoSession').returns(false) + await assert.rejects(async () => await auth.getToken()) + }) + }) + + describe('getTelemetryMetadata', function () { + it('returns valid metadata for BuilderId connection', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + const metadata = await auth.getTelemetryMetadata() + assert.strictEqual(metadata.credentialSourceId, 'awsId') + assert.strictEqual(metadata.credentialStartUrl, constants.builderIdStartUrl) + }) + + it('returns valid metadata for IDC connection', async function () { + await auth.login('https://example.awsapps.com/start', 'us-east-1') + const metadata = await auth.getTelemetryMetadata() + assert.strictEqual(metadata.credentialSourceId, 'iamIdentityCenter') + assert.strictEqual(metadata.credentialStartUrl, 'https://example.awsapps.com/start') + }) + + it('returns undefined metadata when not connected', async function () { + await auth.logout() + const metadata = await auth.getTelemetryMetadata() + assert.strictEqual(metadata.id, 'undefined') + }) + }) + + describe('getAuthFormIds', function () { + it('returns empty array when not connected', async function () { + await auth.logout() + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, []) + }) + + it('returns BuilderId forms when using BuilderId', async function () { + await auth.login(constants.builderIdStartUrl, constants.builderIdRegion) + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, ['builderIdCodeWhisperer']) + }) + + it('returns IDC forms when using IDC without SSO account access', async function () { + const session = (auth as any).session + sinon.stub(session, 'getProfile').resolves({ + ssoSession: { + settings: { + sso_registration_scopes: ['codewhisperer:*'], + }, + }, + }) + + await auth.login('https://example.awsapps.com/start', 'us-east-1') + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, ['identityCenterCodeWhisperer']) + }) + + it('returns IDC forms with explorer when using IDC with SSO account access', async function () { + const session = (auth as any).session + sinon.stub(session, 'getProfile').resolves({ + ssoSession: { + settings: { + sso_registration_scopes: ['codewhisperer:*', 'sso:account:access'], + }, + }, + }) + + await auth.login('https://example.awsapps.com/start', 'us-east-1') + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms.sort(), ['identityCenterCodeWhisperer', 'identityCenterExplorer'].sort()) + }) + + it('returns credentials form for IAM credentials', async function () { + sinon.stub(auth, 'isSsoSession').returns(false) + sinon.stub(auth, 'isConnected').returns(true) + + const forms = await auth.getAuthFormIds() + assert.deepStrictEqual(forms, ['credentials']) + }) + }) + + describe('stateChangeHandler', function () { + let mockLspAuth: any + let regionProfileManager: any + + beforeEach(function () { + mockLspAuth = (auth as any).lspAuth + regionProfileManager = (auth as any).regionProfileManager + }) + + it('updates bearer token when state is refreshed', async function () { + await auth.login(constants.builderIdStartUrl, 'us-east-1') + + await (auth as any).stateChangeHandler({ state: 'refreshed' }) + + assert.ok(mockLspAuth.updateBearerToken.called) + assert.strictEqual(mockLspAuth.updateBearerToken.firstCall.args[0].data, 'fake-data') + }) + + it('cleans up when connection expires', async function () { + await auth.login(constants.builderIdStartUrl, 'us-east-1') + + await (auth as any).stateChangeHandler({ state: 'expired' }) + + assert.ok(mockLspAuth.deleteBearerToken.called) + }) + + it('deletes bearer token when disconnected', async function () { + await (auth as any).stateChangeHandler({ state: 'notConnected' }) + + assert.ok(mockLspAuth.deleteBearerToken.called) + }) + + it('updates bearer token and restores profile on reconnection', async function () { + const restoreProfileSelectionSpy = sinon.spy(regionProfileManager, 'restoreProfileSelection') + + await auth.login('https://example.awsapps.com/start', 'us-east-1') + + await (auth as any).stateChangeHandler({ state: 'connected' }) + + assert.ok(mockLspAuth.updateBearerToken.called) + assert.ok(restoreProfileSelectionSpy.called) + }) + + it('clears region profile cache and invalidates profile on IDC connection expiration', async function () { + const invalidateProfileSpy = sinon.spy(regionProfileManager, 'invalidateProfile') + const clearCacheSpy = sinon.spy(regionProfileManager, 'clearCache') + + await auth.login('https://example.awsapps.com/start', 'us-east-1') + + await (auth as any).stateChangeHandler({ state: 'expired' }) + + assert.ok(invalidateProfileSpy.called) + assert.ok(clearCacheSpy.called) + }) + }) +}) From 4191f6538982f290d8db89b35f7315f30fced046 Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 8 May 2025 15:11:22 -0400 Subject: [PATCH 2/7] Use the testAuthUtil in all unit tests --- .../amazonq/test/unit/amazonq/backend_amazonq.test.ts | 7 +++---- .../codewhisperer/commands/invokeRecommendation.test.ts | 3 ++- .../codewhisperer/region/regionProfileManager.test.ts | 5 ++--- .../codewhisperer/service/inlineCompletionService.test.ts | 8 +++++++- .../unit/codewhisperer/service/keyStrokeHandler.test.ts | 4 ++++ .../amazonq/test/unit/codewhisperer/util/authUtil.test.ts | 8 +++----- 6 files changed, 21 insertions(+), 14 deletions(-) diff --git a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts index 205c19ad798..7db5f2bd704 100644 --- a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts +++ b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts @@ -16,9 +16,8 @@ describe('Amazon Q Login', async function () { let sandbox: sinon.SinonSandbox let backend: backendAmazonQ.AmazonQLoginWebview - await createTestAuthUtil() - - beforeEach(function () { + beforeEach(async function () { + await createTestAuthUtil() sandbox = sinon.createSandbox() backend = new backendAmazonQ.AmazonQLoginWebview() }) @@ -100,7 +99,7 @@ describe('Amazon Q Login', async function () { }) }) - it('signs out of reauth and emits telemetry', async function () { + it.skip('signs out of reauth and emits telemetry', async function () { await backend.signout() assert.ok(!AuthUtil.instance.isConnected()) diff --git a/packages/amazonq/test/unit/codewhisperer/commands/invokeRecommendation.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/invokeRecommendation.test.ts index 68cebe37bb1..56f72edfd3f 100644 --- a/packages/amazonq/test/unit/codewhisperer/commands/invokeRecommendation.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/commands/invokeRecommendation.test.ts @@ -5,7 +5,7 @@ import assert from 'assert' import * as sinon from 'sinon' -import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test' +import { resetCodeWhispererGlobalVariables, createMockTextEditor, createTestAuthUtil } from 'aws-core-vscode/test' import { ConfigurationEntry, invokeRecommendation, @@ -20,6 +20,7 @@ describe('invokeRecommendation', function () { let mockClient: DefaultCodeWhispererClient beforeEach(async function () { + await createTestAuthUtil() await resetCodeWhispererGlobalVariables() getRecommendationStub = sinon.stub(InlineCompletionService.instance, 'getPaginatedRecommendation') }) diff --git a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts index b60d9985eb3..ba4001e5a68 100644 --- a/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts @@ -23,8 +23,6 @@ describe('RegionProfileManager', async function () { description: 'foo description', } - await createTestAuthUtil() - async function setupConnection(type: 'builderId' | 'idc') { if (type === 'builderId') { await AuthUtil.instance.login(constants.builderIdStartUrl, region) @@ -37,7 +35,8 @@ describe('RegionProfileManager', async function () { } } - beforeEach(function () { + beforeEach(async function () { + await createTestAuthUtil() regionProfileManager = new RegionProfileManager(AuthUtil.instance) }) diff --git a/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts b/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts index f24ce9d3f89..dd0bd65505f 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts @@ -19,7 +19,12 @@ import { listCodeWhispererCommandsId, DefaultCodeWhispererClient, } from 'aws-core-vscode/codewhisperer' -import { createMockTextEditor, resetCodeWhispererGlobalVariables, createMockDocument } from 'aws-core-vscode/test' +import { + createMockTextEditor, + resetCodeWhispererGlobalVariables, + createMockDocument, + createTestAuthUtil, +} from 'aws-core-vscode/test' describe('inlineCompletionService', function () { beforeEach(async function () { @@ -192,6 +197,7 @@ describe('codewhisperer status bar', function () { } beforeEach(async function () { + await createTestAuthUtil() await resetCodeWhispererGlobalVariables() sandbox = sinon.createSandbox() statusBar = new TestStatusBar() diff --git a/packages/amazonq/test/unit/codewhisperer/service/keyStrokeHandler.test.ts b/packages/amazonq/test/unit/codewhisperer/service/keyStrokeHandler.test.ts index 4b6a5291f22..f3fa7b399d1 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/keyStrokeHandler.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/keyStrokeHandler.test.ts @@ -9,6 +9,7 @@ import * as sinon from 'sinon' import * as codewhispererSdkClient from 'aws-core-vscode/codewhisperer' import { createMockTextEditor, + createTestAuthUtil, createTextDocumentChangeEvent, resetCodeWhispererGlobalVariables, } from 'aws-core-vscode/test' @@ -160,13 +161,16 @@ describe('keyStrokeHandler', function () { describe('invokeAutomatedTrigger', function () { let mockClient: codewhispererSdkClient.DefaultCodeWhispererClient + beforeEach(async function () { + await createTestAuthUtil() sinon.restore() mockClient = new codewhispererSdkClient.DefaultCodeWhispererClient() await resetCodeWhispererGlobalVariables() sinon.stub(mockClient, 'listRecommendations') sinon.stub(mockClient, 'generateRecommendations') }) + afterEach(function () { sinon.restore() }) diff --git a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts index 8ed2188618b..fa2956f5cf7 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts @@ -11,13 +11,11 @@ import { constants } from 'aws-core-vscode/auth' import { auth2 } from 'aws-core-vscode/auth' describe('AuthUtil', async function () { - await createTestAuthUtil() - - const auth = AuthUtil.instance - const originalSession = (auth as any).session + let auth: any beforeEach(async function () { - ;(auth as any).session = originalSession + await createTestAuthUtil() + auth = AuthUtil.instance }) afterEach(async function () { From d9626834b5251af2a5389438f0c404c94b2e8f08 Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 8 May 2025 15:12:30 -0400 Subject: [PATCH 3/7] Use testAuthUtil in all core tests --- .../test/unit/amazonq/backend_amazonq.test.ts | 2 +- packages/core/src/codewhisperer/util/authUtil.ts | 4 ++++ .../core/src/test/amazonq/customizationUtil.test.ts | 8 +++----- .../src/test/codewhisperer/startSecurityScan.test.ts | 12 +++--------- .../test/codewhispererChat/editor/codelens.test.ts | 8 +++----- packages/core/src/test/shared/featureConfig.test.ts | 11 ++++------- packages/core/src/test/testAuthUtil.ts | 3 +++ 7 files changed, 21 insertions(+), 27 deletions(-) diff --git a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts index 7db5f2bd704..a6711471687 100644 --- a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts +++ b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts @@ -26,7 +26,7 @@ describe('Amazon Q Login', async function () { sandbox.restore() }) - it('signs into builder ID and emits telemetry', async function () { + it('signs into builder ID npm run and emits telemetry', async function () { await backend.startBuilderIdSetup() assert.ok(AuthUtil.instance.isConnected()) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index 092b5f88d7c..a25008796a7 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -82,6 +82,10 @@ export class AuthUtil implements IAuthProvider { }) } + static destroy(): void { + this.#instance = undefined as any + } + isSsoSession() { return this.session.loginType === LoginTypes.SSO } diff --git a/packages/core/src/test/amazonq/customizationUtil.test.ts b/packages/core/src/test/amazonq/customizationUtil.test.ts index bb852afba2e..7e9e209b31f 100644 --- a/packages/core/src/test/amazonq/customizationUtil.test.ts +++ b/packages/core/src/test/amazonq/customizationUtil.test.ts @@ -19,7 +19,7 @@ import { import { FeatureContext, globals } from '../../shared' import { resetCodeWhispererGlobalVariables } from '../codewhisperer/testUtil' import { createSsoProfile, createTestAuth } from '../credentials/testUtil' -import { LanguageClientAuth } from '../../auth/auth2' +import { createTestAuthUtil } from '../testAuthUtil' const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start' @@ -30,13 +30,11 @@ describe('CodeWhisperer-customizationUtils', function () { before(async function () { createTestAuth(globals.globalState) tryRegister(refreshStatusBar) - const mockLspAuth: Partial = { - registerSsoTokenChangedHandler: sinon.stub().resolves(), - } - AuthUtil.create(mockLspAuth as LanguageClientAuth) }) beforeEach(async function () { + await createTestAuthUtil() + auth = createTestAuth(globals.globalState) await auth.createInvalidSsoConnection( createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes }) diff --git a/packages/core/src/test/codewhisperer/startSecurityScan.test.ts b/packages/core/src/test/codewhisperer/startSecurityScan.test.ts index 18a2235d3ee..38b00a2bdd3 100644 --- a/packages/core/src/test/codewhisperer/startSecurityScan.test.ts +++ b/packages/core/src/test/codewhisperer/startSecurityScan.test.ts @@ -28,9 +28,9 @@ import { import * as model from '../../codewhisperer/models/model' import * as errors from '../../shared/errors' import * as timeoutUtils from '../../shared/utilities/timeoutUtils' -import { AuthUtil, SecurityIssueTreeViewProvider } from '../../codewhisperer' +import { SecurityIssueTreeViewProvider } from '../../codewhisperer' import { createClient, mockGetCodeScanResponse } from './testUtil' -import { LanguageClientAuth } from '../../auth/auth2' +import { createTestAuthUtil } from '../testAuthUtil' let extensionContext: FakeExtensionContext let mockSecurityPanelViewProvider: SecurityPanelViewProvider @@ -42,14 +42,8 @@ let focusStub: sinon.SinonStub describe('startSecurityScan', function () { const workspaceFolder = getTestWorkspaceFolder() - before(async function () { - const mockLspAuth: Partial = { - registerSsoTokenChangedHandler: sinon.stub().resolves(), - } - AuthUtil.create(mockLspAuth as LanguageClientAuth) - }) - beforeEach(async function () { + await createTestAuthUtil() extensionContext = await FakeExtensionContext.create() mockSecurityPanelViewProvider = new SecurityPanelViewProvider(extensionContext) appRoot = join(workspaceFolder, 'python3.7-plain-sam-app') diff --git a/packages/core/src/test/codewhispererChat/editor/codelens.test.ts b/packages/core/src/test/codewhispererChat/editor/codelens.test.ts index 8ae1d715004..3e9bd9284a5 100644 --- a/packages/core/src/test/codewhispererChat/editor/codelens.test.ts +++ b/packages/core/src/test/codewhispererChat/editor/codelens.test.ts @@ -24,7 +24,8 @@ import { PressTabState, TryMoreExState, } from '../../../codewhisperer/views/lineAnnotationController' -import { AuthState, LanguageClientAuth } from '../../../auth/auth2' +import { AuthState } from '../../../auth/auth2' +import { createTestAuthUtil } from '../../testAuthUtil' describe('TryChatCodeLensProvider', () => { let instance: TryChatCodeLensProvider @@ -41,13 +42,10 @@ describe('TryChatCodeLensProvider', () => { // that originally would have been registered by the `core` `activate()` at some point tryRegister(tryChatCodeLensCommand) tryRegister(focusAmazonQPanel) - const mockLspAuth: Partial = { - registerSsoTokenChangedHandler: sinon.stub().resolves(), - } - AuthUtil.create(mockLspAuth as LanguageClientAuth) }) beforeEach(async function () { + await createTestAuthUtil() isAmazonQVisibleEventEmitter = new vscode.EventEmitter() isAmazonQVisibleEvent = isAmazonQVisibleEventEmitter.event instance = new TryChatCodeLensProvider(isAmazonQVisibleEvent, () => codeLensPosition) diff --git a/packages/core/src/test/shared/featureConfig.test.ts b/packages/core/src/test/shared/featureConfig.test.ts index dd84f59d554..e94358361b0 100644 --- a/packages/core/src/test/shared/featureConfig.test.ts +++ b/packages/core/src/test/shared/featureConfig.test.ts @@ -7,18 +7,15 @@ import assert from 'assert' import sinon from 'sinon' import { AWSError, Request } from 'aws-sdk' import { Features, FeatureConfigProvider, featureDefinitions, FeatureName } from '../../shared/featureConfig' -import { AuthUtil, ListFeatureEvaluationsResponse } from '../../codewhisperer' +import { ListFeatureEvaluationsResponse } from '../../codewhisperer' import { createSpyClient } from '../codewhisperer/testUtil' import { mockFeatureConfigsData } from '../fake/mockFeatureConfigData' -import { LanguageClientAuth } from '../../auth/auth2' +import { createTestAuthUtil } from '../testAuthUtil' describe('FeatureConfigProvider', () => { - const mockLspAuth: Partial = { - registerSsoTokenChangedHandler: sinon.stub().resolves(), - } - AuthUtil.create(mockLspAuth as LanguageClientAuth) - beforeEach(async () => { + await createTestAuthUtil() + const clientSpy = await createSpyClient() sinon.stub(clientSpy, 'listFeatureEvaluations').returns({ promise: () => diff --git a/packages/core/src/test/testAuthUtil.ts b/packages/core/src/test/testAuthUtil.ts index 14e7045739a..591eb52d270 100644 --- a/packages/core/src/test/testAuthUtil.ts +++ b/packages/core/src/test/testAuthUtil.ts @@ -10,6 +10,9 @@ import { LanguageClientAuth } from '../auth/auth2' import { AuthUtil } from '../codewhisperer/util/authUtil' export async function createTestAuthUtil() { + sinon.restore() + AuthUtil.destroy() + const encryptionKey = crypto.randomBytes(32) const jwe = await new jose.CompactEncrypt(new TextEncoder().encode(JSON.stringify({ your: 'mock data' }))) From 47e126f5e3c06581159f2cc1e33019688b864439 Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 8 May 2025 15:44:20 -0400 Subject: [PATCH 4/7] Re-enable test --- packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts index a6711471687..d69dcac39de 100644 --- a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts +++ b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts @@ -26,7 +26,7 @@ describe('Amazon Q Login', async function () { sandbox.restore() }) - it('signs into builder ID npm run and emits telemetry', async function () { + it('signs into builder ID and emits telemetry', async function () { await backend.startBuilderIdSetup() assert.ok(AuthUtil.instance.isConnected()) @@ -99,7 +99,7 @@ describe('Amazon Q Login', async function () { }) }) - it.skip('signs out of reauth and emits telemetry', async function () { + it('signs out of reauth and emits telemetry', async function () { await backend.signout() assert.ok(!AuthUtil.instance.isConnected()) From e9c2d40c2f3568464fa8a628b95ca03d92554498 Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 8 May 2025 15:55:49 -0400 Subject: [PATCH 5/7] Fix failing unit test --- packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts | 4 +++- packages/core/src/codewhisperer/util/authUtil.ts | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts index d69dcac39de..f774c09de44 100644 --- a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts +++ b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts @@ -99,7 +99,9 @@ describe('Amazon Q Login', async function () { }) }) - it('signs out of reauth and emits telemetry', async function () { + it.only('signs out of reauth and emits telemetry', async function () { + await getStartUrl.connectToEnterpriseSso(startUrl, region) + await backend.signout() assert.ok(!AuthUtil.instance.isConnected()) diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts index a25008796a7..362b1ae7157 100644 --- a/packages/core/src/codewhisperer/util/authUtil.ts +++ b/packages/core/src/codewhisperer/util/authUtil.ts @@ -82,6 +82,7 @@ export class AuthUtil implements IAuthProvider { }) } + // Do NOT use this in production code, only used for testing static destroy(): void { this.#instance = undefined as any } From dc8340578146a6bdfae85727605aac555be76694 Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 8 May 2025 16:05:16 -0400 Subject: [PATCH 6/7] Linter --- packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts index f774c09de44..5d9972019f4 100644 --- a/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts +++ b/packages/amazonq/test/unit/amazonq/backend_amazonq.test.ts @@ -99,7 +99,7 @@ describe('Amazon Q Login', async function () { }) }) - it.only('signs out of reauth and emits telemetry', async function () { + it('signs out of reauth and emits telemetry', async function () { await getStartUrl.connectToEnterpriseSso(startUrl, region) await backend.signout() From 86206f22ec65ca88b4e069c27a2c0d20bb675cc9 Mon Sep 17 00:00:00 2001 From: opieter-aws Date: Thu, 8 May 2025 16:40:40 -0400 Subject: [PATCH 7/7] Address comment --- packages/core/src/test/testAuthUtil.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/core/src/test/testAuthUtil.ts b/packages/core/src/test/testAuthUtil.ts index 591eb52d270..4feefec2d68 100644 --- a/packages/core/src/test/testAuthUtil.ts +++ b/packages/core/src/test/testAuthUtil.ts @@ -10,9 +10,6 @@ import { LanguageClientAuth } from '../auth/auth2' import { AuthUtil } from '../codewhisperer/util/authUtil' export async function createTestAuthUtil() { - sinon.restore() - AuthUtil.destroy() - const encryptionKey = crypto.randomBytes(32) const jwe = await new jose.CompactEncrypt(new TextEncoder().encode(JSON.stringify({ your: 'mock data' }))) @@ -42,5 +39,8 @@ export async function createTestAuthUtil() { encryptionKey, } + // Since AuthUtil is a singleton, we want to remove an existing instance before setting up a new one + AuthUtil.destroy() + AuthUtil.create(mockLspAuth as LanguageClientAuth) }