Skip to content

Commit 4489ec6

Browse files
justinmk3justinmk3
committed
refactor(auth): use globalState abstraction
1 parent d5d026a commit 4489ec6

File tree

14 files changed

+81
-53
lines changed

14 files changed

+81
-53
lines changed

‎packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import {
2121
createTestAuth,
2222
} from 'aws-core-vscode/test'
2323
import { Auth, Connection, isAnySsoConnection, isBuilderIdConnection } from 'aws-core-vscode/auth'
24-
import { vscodeComponent } from 'aws-core-vscode/shared'
24+
import { globals, vscodeComponent } from 'aws-core-vscode/shared'
2525

2626
const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
2727

@@ -30,7 +30,7 @@ describe('AuthUtil', async function () {
3030
let authUtil: AuthUtil
3131

3232
beforeEach(async function () {
33-
auth = createTestAuth()
33+
auth = createTestAuth(globals.globalState)
3434
authUtil = new AuthUtil(auth)
3535
})
3636

@@ -286,7 +286,7 @@ describe('getChatAuthState()', function () {
286286
let laterDate: Date
287287

288288
beforeEach(async function () {
289-
auth = createTestAuth()
289+
auth = createTestAuth(globals.globalState)
290290
authUtil = new AuthUtil(auth)
291291

292292
laterDate = new Date(Date.now() + 10_000_000)

‎packages/core/src/auth/secondaryAuth.ts‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -98,8 +98,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
9898
public readonly toolId: ToolId,
9999
public readonly toolLabel: string,
100100
public readonly isUsable: (conn: Connection) => conn is T,
101-
private readonly auth: Auth,
102-
private readonly memento = globals.context.globalState
101+
private readonly auth: Auth
103102
) {
104103
const handleConnectionChanged = async (newActiveConn?: Connection) => {
105104
if (newActiveConn === undefined && this.#activeConnection?.id) {
@@ -176,7 +175,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
176175
}
177176

178177
public async saveConnection(conn: T) {
179-
await this.memento.update(this.key, conn.id)
178+
await globals.context.globalState.update(this.key, conn.id)
180179
this.#savedConnection = conn
181180
this.#onDidChangeActiveConnection.fire(this.activeConnection)
182181
}
@@ -207,7 +206,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
207206

208207
/** Stop using the saved connection and fallback to using the active connection, if it is usable. */
209208
public async clearSavedConnection() {
210-
await this.memento.update(this.key, undefined)
209+
await globals.context.globalState.update(this.key, undefined)
211210
this.#savedConnection = undefined
212211
this.#onDidChangeActiveConnection.fire(this.activeConnection)
213212
}
@@ -252,18 +251,18 @@ export class SecondaryAuth<T extends Connection = Connection> {
252251
})
253252

254253
private async loadSavedConnection() {
255-
const id = cast(this.memento.get(this.key), Optional(String))
254+
const id = cast(globals.context.globalState.get(this.key), Optional(String))
256255
if (id === undefined) {
257256
return
258257
}
259258

260259
const conn = await this.auth.getConnection({ id })
261260
if (conn === undefined) {
262261
getLogger().warn(`auth (${this.toolId}): removing saved connection "${this.key}" as it no longer exists`)
263-
await this.memento.update(this.key, undefined)
262+
await globals.context.globalState.update(this.key, undefined)
264263
} else if (!this.isUsable(conn)) {
265264
getLogger().warn(`auth (${this.toolId}): saved connection "${this.key}" is not valid`)
266-
await this.memento.update(this.key, undefined)
265+
await globals.context.globalState.update(this.key, undefined)
267266
} else {
268267
await this.auth.refreshConnectionState(conn)
269268
return conn

‎packages/core/src/shared/credentials/credentialsProfileMru.ts‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,21 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
import * as vscode from 'vscode'
6+
import globals from '../extensionGlobals'
77

88
/**
99
* Tracks the credentials selected by the user, ordered by most recent.
1010
*/
1111
export class CredentialsProfileMru {
1212
public static readonly maxCredentialMruSize = 5
1313

14-
private static readonly configurationStateName: string = 'recentCredentials'
15-
16-
public constructor(private readonly _context: vscode.ExtensionContext) {}
14+
public constructor() {}
1715

1816
/**
1917
* @description Returns the most recently used credentials names
2018
*/
2119
public getMruList(): string[] {
22-
return this._context.globalState.get<string[]>(CredentialsProfileMru.configurationStateName, [])
20+
return globals.globalState.tryGet<string[]>('recentCredentials', Object, [])
2321
}
2422

2523
/**
@@ -38,6 +36,6 @@ export class CredentialsProfileMru {
3836

3937
mru.splice(CredentialsProfileMru.maxCredentialMruSize)
4038

41-
await this._context.globalState.update(CredentialsProfileMru.configurationStateName, mru)
39+
await globals.globalState.update('recentCredentials', mru)
4240
}
4341
}

‎packages/core/src/shared/credentials/defaultCredentialSelectionDataProvider.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ export class DefaultCredentialSelectionDataProvider implements CredentialSelecti
4646
public readonly existingProfileNames: string[],
4747
protected context: vscode.ExtensionContext
4848
) {
49-
this._credentialsMru = new CredentialsProfileMru(context)
49+
this._credentialsMru = new CredentialsProfileMru()
5050
}
5151

5252
public async pickCredentialProfile(

‎packages/core/src/shared/utilities/mementos.ts‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,14 @@ import globals from '../extensionGlobals'
99
import { getCodeCatalystDevEnvId } from '../vscode/env'
1010

1111
/**
12-
* Divides a memento at the specified key, creating a 'scoped' memento.
12+
* Creates a memento interface to a nested object stored at `key`.
13+
*
14+
* For example, `partition(m, 'foo')` creates a nested object at "foo":
15+
*
16+
* "foo": {
17+
* }
18+
*
19+
* and returns a memento that gets/sets keys only on "foo" (not its container `m`).
1320
*/
1421
export function partition(memento: vscode.Memento, key: string): vscode.Memento {
1522
const get = () => cast(memento.get(key), Optional(Record(String, Unknown)))
@@ -22,19 +29,20 @@ export function partition(memento: vscode.Memento, key: string): vscode.Memento
2229
}
2330
}
2431

32+
/** Avoids sharing globalState with remote vscode instances. */
2533
export function getEnvironmentSpecificMemento(): vscode.Memento {
2634
if (!vscode.env.remoteName) {
2735
// local compute: no further partitioning
28-
return globals.context.globalState
36+
return globals.globalState
2937
}
3038

3139
const devEnvId = getCodeCatalystDevEnvId()
3240

3341
if (devEnvId !== undefined) {
3442
// dev env: partition to dev env ID (compute backend might not always be the same)
35-
return partition(globals.context.globalState, devEnvId)
43+
return partition(globals.globalState, devEnvId)
3644
}
3745

3846
// remote env: keeps a shared "global state" for all workspaces that report the same machine ID
39-
return partition(globals.context.globalState, globals.machineId)
47+
return partition(globals.globalState, globals.machineId)
4048
}

‎packages/core/src/test/codecatalyst/auth.test.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { FakeSecretStorage } from '../fakeExtensionContext'
1010
import { createBuilderIdProfile, createSsoProfile, createTestAuth } from '../credentials/testUtil'
1111
import Sinon from 'sinon'
1212
import { isAnySsoConnection } from '../../auth/connection'
13+
import globals from '../../shared/extensionGlobals'
1314

1415
const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
1516

@@ -18,7 +19,7 @@ describe('CodeCatalystAuthenticationProvider', async function () {
1819
let codecatalystAuth: CodeCatalystAuthenticationProvider
1920

2021
beforeEach(async function () {
21-
auth = createTestAuth()
22+
auth = createTestAuth(globals.globalState)
2223
codecatalystAuth = new CodeCatalystAuthenticationProvider(
2324
new CodeCatalystAuthStorage(new FakeSecretStorage()),
2425
auth

‎packages/core/src/test/credentials/auth.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ describe('Auth', function () {
2828
let auth: ReturnType<typeof createTestAuth>
2929

3030
beforeEach(function () {
31-
auth = createTestAuth()
31+
auth = createTestAuth(globals.globalState)
3232
})
3333

3434
it('can create a new sso connection', async function () {

‎packages/core/src/test/credentials/secondaryAuth.test.ts‎

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,33 +21,33 @@ describe('SecondaryAuth', function () {
2121
let onDidChangeActiveConnection: SinonStub
2222

2323
beforeEach(async function () {
24-
auth = createTestAuth()
24+
auth = createTestAuth(globals.globalState)
2525
sandbox = createSandbox()
2626
conn = await auth.createConnection(createBuilderIdProfile({ scopes: scopes }))
2727
isValid = (conn: Connection): conn is Connection => {
2828
return isSsoConnection(conn) && hasScopes(conn, scopes)
2929
}
30-
secondaryAuth = getSecondaryAuth(auth, 'testId', 'testLabel', isValid)
30+
// await globals.globalState.clear()
31+
secondaryAuth = getSecondaryAuth(auth, 'codecatalyst', 'testLabel', isValid)
3132
onDidChangeActiveConnection = sandbox.stub()
3233
secondaryAuth.onDidChangeActiveConnection(onDidChangeActiveConnection)
33-
await globals.globalState.clear()
3434
})
3535

3636
afterEach(async function () {
3737
sandbox.restore()
3838
})
3939

40-
it('When no SecondaryAuth set or valid PrimaryAuth exist', async function () {
40+
it('no SecondaryAuth set or valid PrimaryAuth exist', async function () {
4141
assert.strictEqual(secondaryAuth.activeConnection?.id, undefined)
4242
})
4343

44-
it('When no SecondaryAuth set but valid PrimaryAuth is used', async function () {
44+
it('no SecondaryAuth set but valid PrimaryAuth is used', async function () {
4545
await auth.useConnection(conn)
4646
assert.strictEqual(onDidChangeActiveConnection.calledOnce, true)
4747
assert.strictEqual(secondaryAuth.activeConnection?.id, conn.id)
4848
})
4949

50-
it('When valid PrimaryAuth is set BUT SecondaryAuth is already using the same connection', async function () {
50+
it('valid PrimaryAuth is set BUT SecondaryAuth is already using the same connection', async function () {
5151
await secondaryAuth.useNewConnection(conn)
5252
// we save this connection so we expect it to trigger an event
5353
assert.strictEqual(onDidChangeActiveConnection.called, true)
@@ -60,7 +60,7 @@ describe('SecondaryAuth', function () {
6060
assert.strictEqual(secondaryAuth.activeConnection?.id, conn.id)
6161
})
6262

63-
it('When valid PrimaryAuth changes BUT no SecondaryAuth is set + valid PrimaryAuth currently exists', async function () {
63+
it('valid PrimaryAuth changes BUT SecondaryAuth not set + valid PrimaryAuth exists', async function () {
6464
// Make primary auth already exist
6565
await auth.useConnection(conn)
6666
assert.strictEqual(onDidChangeActiveConnection.called, true)
@@ -74,7 +74,7 @@ describe('SecondaryAuth', function () {
7474
assert.strictEqual(secondaryAuth.activeConnection?.id, otherValidConn.id)
7575
})
7676

77-
it('When valid PrimaryAuth deleted BUT SecondaryAuth is already set', async function () {
77+
it('valid PrimaryAuth deleted BUT SecondaryAuth is set', async function () {
7878
await secondaryAuth.useNewConnection(conn)
7979

8080
// add valid connection to the PrimaryAuth
@@ -89,7 +89,7 @@ describe('SecondaryAuth', function () {
8989
assert.strictEqual(secondaryAuth.activeConnection?.id, conn.id)
9090
})
9191

92-
it('When valid SecondaryAuth deleted BUT valid PrimaryAuth already exists', async function () {
92+
it('valid SecondaryAuth deleted BUT valid PrimaryAuth exists', async function () {
9393
await secondaryAuth.useNewConnection(conn)
9494

9595
// add valid connection to the PrimaryAuth
@@ -101,11 +101,21 @@ describe('SecondaryAuth', function () {
101101
await auth.deleteConnection(conn)
102102

103103
// we fallback to the PrimaryAuth connection
104-
assert.strictEqual(secondaryAuth.activeConnection?.id, otherConn.id)
105104
assert.strictEqual(onDidChangeActiveConnection.called, true)
105+
assert.deepStrictEqual(
106+
{
107+
id: secondaryAuth.activeConnection?.id,
108+
label: secondaryAuth.activeConnection?.label,
109+
},
110+
{
111+
id: otherConn.id,
112+
label: otherConn.label,
113+
}
114+
)
115+
assert.strictEqual(secondaryAuth.activeConnection?.id, otherConn.id)
106116
})
107117

108-
it('When SecondaryAuth is invalid, but is reauthenticated elsewhere', async function () {
118+
it('SecondaryAuth is invalid, but is reauthenticated elsewhere', async function () {
109119
// PrimaryAuth has its own conn, we don't care about this
110120
await auth.useConnection(conn)
111121

‎packages/core/src/test/credentials/testUtil.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import { CredentialsProviderManager } from '../../auth/providers/credentialsProv
1010
import { SsoClient } from '../../auth/sso/clients'
1111
import { builderIdStartUrl, SsoToken } from '../../auth/sso/model'
1212
import { DeviceFlowAuthorization, SsoAccessTokenProvider } from '../../auth/sso/ssoAccessTokenProvider'
13-
import { FakeMemento } from '../fakeExtensionContext'
1413
import { captureEvent, EventCapturer } from '../testUtil'
1514
import { stub } from '../utilities/stubber'
1615
import globals from '../../shared/extensionGlobals'
@@ -29,8 +28,17 @@ export const ssoConnection: SsoConnection = {
2928
startUrl: 'https://nkomonen.awsapps.com/start',
3029
getToken: sinon.stub(),
3130
}
32-
export const builderIdConnection: SsoConnection = { ...ssoConnection, startUrl: builderIdStartUrl, label: 'builderId' }
33-
export const iamConnection: IamConnection = { type: 'iam', id: '0', label: 'iam', getCredentials: sinon.stub() }
31+
export const builderIdConnection: SsoConnection = {
32+
...ssoConnection,
33+
startUrl: builderIdStartUrl,
34+
label: 'builderId',
35+
}
36+
export const iamConnection: IamConnection = {
37+
type: 'iam',
38+
id: '0',
39+
label: 'iam',
40+
getCredentials: sinon.stub(),
41+
}
3442

3543
export function createSsoProfile(props?: Partial<Omit<SsoProfile, 'type'>>): SsoProfile {
3644
return {
@@ -80,7 +88,7 @@ type TestAuth = Auth & {
8088
getTestTokenProvider(connection: Pick<Connection, 'id'>): ReturnType<typeof createTestTokenProvider>
8189
}
8290

83-
export function createTestAuth(): TestAuth {
91+
export function createTestAuth(globalState: vscode.Memento): TestAuth {
8492
const tokenProviders = new Map<string, ReturnType<typeof createTestTokenProvider>>()
8593

8694
function getTokenProvider(...[profile]: ConstructorParameters<typeof SsoAccessTokenProvider>) {
@@ -108,7 +116,7 @@ export function createTestAuth(): TestAuth {
108116
const ssoClient = stub(SsoClient, { region: 'not set' })
109117
ssoClient.logout.resolves()
110118

111-
const store = new ProfileStore(new FakeMemento())
119+
const store = new ProfileStore(globalState)
112120
const credentialsManager = new CredentialsProviderManager()
113121
const auth = new Auth(store, credentialsManager, () => ssoClient, getTokenProvider)
114122

‎packages/core/src/test/globalSetup.test.ts‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import { CodelensRootRegistry } from '../shared/fs/codelensRootRegistry'
1616
import { CloudFormationTemplateRegistry } from '../shared/fs/templateRegistry'
1717
import { getLogger, LogLevel } from '../shared/logger'
1818
import { setLogger } from '../shared/logger/logger'
19-
import { FakeExtensionContext, FakeMemento } from './fakeExtensionContext'
19+
import { FakeExtensionContext } from './fakeExtensionContext'
2020
import { TestLogger } from './testLogger'
2121
import * as testUtil from './testUtil'
2222
import { getTestWindow, resetTestWindow } from './shared/vscode/window'
@@ -52,7 +52,10 @@ export async function mochaGlobalSetup(extensionId: string) {
5252
const fakeContext = await FakeExtensionContext.create()
5353
fakeContext.globalStorageUri = (await testUtil.createTestWorkspaceFolder('globalStoragePath')).uri
5454
fakeContext.extensionPath = globals.context.extensionPath
55-
Object.assign(globals, { context: fakeContext })
55+
Object.assign(globals, {
56+
context: fakeContext,
57+
globalState: new GlobalState(fakeContext.globalState),
58+
})
5659
}
5760
}
5861

@@ -84,9 +87,8 @@ export const mochaHooks = {
8487
globals.telemetry.clearRecords()
8588
globals.telemetry.logger.clear()
8689
TelemetryDebounceInfo.instance.clear()
87-
const fakeGlobalState = new FakeMemento()
88-
;(globals.context as FakeExtensionContext).globalState = fakeGlobalState
89-
;(globals as any).globalState = new GlobalState(fakeGlobalState)
90+
// mochaGlobalSetup() set this to a fake, so it's safe to clear it here.
91+
await globals.globalState.clear()
9092

9193
await testUtil.closeAllEditors()
9294
},

0 commit comments

Comments
 (0)