refactor(core): use globalState abstraction

justinmk3 · justinmk3 · commit 9cdb8c4398f7 · 2024-07-18T15:15:24.000-07:00
diff --git a/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts
@@ -37,7 +37,6 @@ describe('onAcceptance', function () {
         it('Should enqueue an event object to tracker', async function () {
             const mockEditor = createMockTextEditor()
             const trackerSpy = sinon.spy(CodeWhispererTracker.prototype, 'enqueue')
-            const extensionContext = await FakeExtensionContext.create()
             const fakeReferences = [
                 {
                     message: '',
@@ -49,22 +48,19 @@ describe('onAcceptance', function () {
                     },
                 },
             ]
-            await onAcceptance(
-                {
-                    editor: mockEditor,
-                    range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
-                    effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
-                    acceptIndex: 0,
-                    recommendation: "print('Hello World!')",
-                    requestId: '',
-                    sessionId: '',
-                    triggerType: 'OnDemand',
-                    completionType: 'Line',
-                    language: 'python',
-                    references: fakeReferences,
-                },
-                extensionContext.globalState
-            )
+            await onAcceptance({
+                editor: mockEditor,
+                range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
+                effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
+                acceptIndex: 0,
+                recommendation: "print('Hello World!')",
+                requestId: '',
+                sessionId: '',
+                triggerType: 'OnDemand',
+                completionType: 'Line',
+                language: 'python',
+                references: fakeReferences,
+            })
             const actualArg = trackerSpy.getCall(0).args[0] as AcceptedSuggestionEntry
             assert.ok(trackerSpy.calledOnce)
             assert.strictEqual(actualArg.originalString, 'def two_sum(nums, target):')
@@ -79,7 +75,7 @@ describe('onAcceptance', function () {
         })
 
         it('Should report telemetry that records this user decision event', async function () {
-            await globals.context.globalState.update('CODEWHISPERER_USER_GROUP', {
+            await globals.globalState.update('CODEWHISPERER_USER_GROUP', {
                 group: UserGroup.Control,
                 version: extensionVersion,
             })
@@ -97,23 +93,19 @@ describe('onAcceptance', function () {
             session.triggerType = 'OnDemand'
             session.setCompletionType(0, session.recommendations[0])
             const assertTelemetry = assertTelemetryCurried('codewhisperer_userDecision')
-            const extensionContext = await FakeExtensionContext.create()
-            await onAcceptance(
-                {
-                    editor: mockEditor,
-                    range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 21)),
-                    effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
-                    acceptIndex: 0,
-                    recommendation: "print('Hello World!')",
-                    requestId: '',
-                    sessionId: '',
-                    triggerType: 'OnDemand',
-                    completionType: 'Line',
-                    language: 'python',
-                    references: undefined,
-                },
-                extensionContext.globalState
-            )
+            await onAcceptance({
+                editor: mockEditor,
+                range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 21)),
+                effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
+                acceptIndex: 0,
+                recommendation: "print('Hello World!')",
+                requestId: '',
+                sessionId: '',
+                triggerType: 'OnDemand',
+                completionType: 'Line',
+                language: 'python',
+                references: undefined,
+            })
             assertTelemetry({
                 codewhispererRequestId: 'test',
                 codewhispererSessionId: 'test',
diff --git a/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts
@@ -57,7 +57,7 @@ describe('onInlineAcceptance', function () {
         })
 
         it('Should report telemetry that records this user decision event', async function () {
-            await globals.context.globalState.update('CODEWHISPERER_USER_GROUP', {
+            await globals.globalState.update('CODEWHISPERER_USER_GROUP', {
                 group: UserGroup.Classifier,
                 version: extensionVersion,
             })
diff --git a/packages/amazonq/test/unit/codewhisperer/service/codewhisperer.test.ts b/packages/amazonq/test/unit/codewhisperer/service/codewhisperer.test.ts
@@ -115,7 +115,7 @@ describe('codewhisperer', async function () {
             ideCategory: 'VSCODE',
             operatingSystem: getOperatingSystem(),
             product: 'CodeWhisperer',
-            clientId: getClientId(globals.context.globalState),
+            clientId: getClientId(globals.globalState),
         }
 
         await codeWhispererClient.sendTelemetryEvent({ telemetryEvent: userTriggerDecisionPayload })
diff --git a/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts b/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts
@@ -114,7 +114,7 @@ describe('recommendationHandler', function () {
         })
 
         it('should call telemetry function that records a CodeWhisperer service invocation', async function () {
-            await globals.context.globalState.update('CODEWHISPERER_USER_GROUP', {
+            await globals.globalState.update('CODEWHISPERER_USER_GROUP', {
                 group: UserGroup.CrossFile,
                 version: extensionVersion,
             })
@@ -166,7 +166,7 @@ describe('recommendationHandler', function () {
         })
 
         it('should call telemetry function that records a Empty userDecision event', async function () {
-            await globals.context.globalState.update('CODEWHISPERER_USER_GROUP', {
+            await globals.globalState.update('CODEWHISPERER_USER_GROUP', {
                 group: UserGroup.CrossFile,
                 version: extensionVersion,
             })
diff --git a/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts b/packages/amazonq/test/unit/codewhisperer/tracker/codewhispererTracker.test.ts
@@ -95,7 +95,7 @@ describe('codewhispererTracker', function () {
         })
 
         it('Should call recordCodewhispererUserModification with suggestion event', async function () {
-            await globals.context.globalState.update('CODEWHISPERER_USER_GROUP', {
+            await globals.globalState.update('CODEWHISPERER_USER_GROUP', {
                 group: UserGroup.CrossFile,
                 version: extensionVersion,
             })
diff --git a/packages/core/src/amazonq/onboardingPage/walkthrough.ts b/packages/core/src/amazonq/onboardingPage/walkthrough.ts
@@ -15,12 +15,8 @@ import vscode from 'vscode'
  * Show the Amazon Q walkthrough one time forever when the user adds an Amazon Q connection.
  * All subsequent calls to this do nothing.
  */
-export async function showAmazonQWalkthroughOnce(
-    state = globals.context.globalState,
-    showWalkthrough = () => openAmazonQWalkthrough.execute()
-) {
-    const hasShownWalkthroughId = 'aws.amazonq.hasShownWalkthrough'
-    const hasShownWalkthrough = state.get(hasShownWalkthroughId, false)
+export async function showAmazonQWalkthroughOnce(showWalkthrough = () => openAmazonQWalkthrough.execute()) {
+    const hasShownWalkthrough = globals.globalState.tryGet('aws.amazonq.hasShownWalkthrough', Boolean, false)
     if (hasShownWalkthrough) {
         return
     }
@@ -30,7 +26,7 @@ export async function showAmazonQWalkthroughOnce(
         return
     }
 
-    await state.update(hasShownWalkthroughId, true)
+    await globals.globalState.update('aws.amazonq.hasShownWalkthrough', true)
     await showWalkthrough()
 }
 
diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts
@@ -337,7 +337,7 @@ export class FeatureDevClient {
                     ideCategory: 'VSCODE',
                     operatingSystem: getOperatingSystem(),
                     product: 'FeatureDev', // Should be the same as in JetBrains
-                    clientId: getClientId(globals.context.globalState),
+                    clientId: getClientId(globals.globalState),
                     ideVersion: extensionVersion,
                 },
             }
diff --git a/packages/core/src/auth/secondaryAuth.ts b/packages/core/src/auth/secondaryAuth.ts
@@ -11,9 +11,12 @@ import { cast, Optional } from '../shared/utilities/typeConstructors'
 import { Auth } from './auth'
 import { once, onceChanged } from '../shared/utilities/functionUtils'
 import { isNonNullable } from '../shared/utilities/tsUtils'
+import { ToolIdStateKey } from '../shared/globalState'
 import { Connection, SsoConnection, StatefulConnection } from './connection'
 import { indent } from '../shared/utilities/textUtilities'
 
+export type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
+
 let currentConn: Auth['activeConnection']
 const auths = new Map<string, SecondaryAuth>()
 const multiConnectionListeners = new WeakMap<Auth, vscode.Disposable>()
@@ -36,7 +39,7 @@ const registerAuthListener = (auth: Auth) => {
 
 export function getSecondaryAuth<T extends Connection>(
     auth: Auth,
-    toolId: string,
+    toolId: ToolId,
     toolLabel: string,
     isValid: (conn: Connection) => conn is T
 ): SecondaryAuth<T> {
@@ -87,16 +90,16 @@ export class SecondaryAuth<T extends Connection = Connection> {
     #savedConnection: T | undefined
     protected static readonly logIfChanged = onceChanged((s: string) => getLogger().info(s))
 
-    private readonly key = `${this.toolId}.savedConnectionId`
+    private readonly key: ToolIdStateKey = `${this.toolId}.savedConnectionId`
     readonly #onDidChangeActiveConnection = new vscode.EventEmitter<T | undefined>()
     public readonly onDidChangeActiveConnection = this.#onDidChangeActiveConnection.event
 
     public constructor(
-        public readonly toolId: string,
+        public readonly toolId: ToolId,
         public readonly toolLabel: string,
         public readonly isUsable: (conn: Connection) => conn is T,
         private readonly auth: Auth,
-        private readonly memento = globals.context.globalState
+        private readonly memento = globals.globalState
     ) {
         const handleConnectionChanged = async (newActiveConn?: Connection) => {
             if (newActiveConn === undefined && this.#activeConnection?.id) {
diff --git a/packages/core/src/auth/utils.ts b/packages/core/src/auth/utils.ts
@@ -629,15 +629,12 @@ export class ExtensionUse {
     /**
      * Check if this is the first use/session of the extension.
      */
-    isFirstUse(
-        state: vscode.Memento = globals.context.globalState,
-        hasExistingConnections = () => Auth.instance.hasConnections
-    ): boolean {
+    isFirstUse(hasExistingConnections = () => Auth.instance.hasConnections): boolean {
         if (this.isFirstUseCurrentSession !== undefined) {
             return this.isFirstUseCurrentSession
         }
 
-        this.isFirstUseCurrentSession = state.get(this.isExtensionFirstUseKey)
+        this.isFirstUseCurrentSession = globals.globalState.get('isExtensionFirstUse')
         if (this.isFirstUseCurrentSession === undefined) {
             // The variable in the store is not defined yet, fallback to checking if they have existing connections.
             this.isFirstUseCurrentSession = !hasExistingConnections()
@@ -650,7 +647,7 @@ export class ExtensionUse {
         }
 
         // Update state, so next time it is not first use
-        this.updateMemento(state, this.isExtensionFirstUseKey, false)
+        this.updateMemento('isExtensionFirstUse', false)
 
         return this.isFirstUseCurrentSession
     }
@@ -663,22 +660,20 @@ export class ExtensionUse {
      * Caveat: This may return true even if an update hadn't occurred IF
      * this function hasn't been called.
      */
-    wasUpdated(state: vscode.Memento = globals.context.globalState) {
+    wasUpdated() {
         if (this.wasExtensionUpdated !== undefined) {
             return this.wasExtensionUpdated
         }
         const currentVersion = extensionVersion
 
-        this.wasExtensionUpdated = currentVersion !== state.get(this.lastExtensionVersionKey)
-        this.updateMemento(state, this.lastExtensionVersionKey, currentVersion)
+        this.wasExtensionUpdated = currentVersion !== globals.globalState.get('lastExtensionVersion')
+        this.updateMemento(this.lastExtensionVersionKey, currentVersion)
 
         return this.wasExtensionUpdated
     }
 
-    private updateMemento(memento: vscode.Memento, key: string, val: any) {
-        memento.update(key, val).then(undefined, (e) => {
-            getLogger().error('Memento.update failed: %s', (e as Error).message)
-        })
+    private updateMemento(key: 'isExtensionFirstUse' | 'lastExtensionVersion', val: any) {
+        globals.globalState.tryUpdate(key, val)
     }
 
     static #instance: ExtensionUse
diff --git a/packages/core/src/codewhisperer/commands/onAcceptance.ts b/packages/core/src/codewhisperer/commands/onAcceptance.ts
@@ -19,7 +19,7 @@ import path from 'path'
 /**
  * This function is called when user accepts a intelliSense suggestion or an inline suggestion
  */
-export async function onAcceptance(acceptanceEntry: OnRecommendationAcceptanceEntry, globalStorage: vscode.Memento) {
+export async function onAcceptance(acceptanceEntry: OnRecommendationAcceptanceEntry) {
     RecommendationHandler.instance.cancelPaginatedRequest()
     /**
      * Format document
diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts
@@ -8,6 +8,9 @@ import { getLogger } from './logger/logger'
 import * as redshift from '../awsService/redshift/models/models'
 import { TypeConstructor, cast } from './utilities/typeConstructors'
 
+type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
+export type ToolIdStateKey = `${ToolId}.savedConnectionId`
+
 type samInitStateKey =
     | 'ACTIVATION_TEMPLATE_PATH_KEY'
     | 'ACTIVATION_LAUNCH_PATH_KEY'
@@ -20,15 +23,17 @@ type stepFunctionsKey = 'SCRIPT_LAST_DOWNLOADED_URL' | 'CSS_LAST_DOWNLOADED_URL'
 type globalKey =
     | samInitStateKey
     | stepFunctionsKey
+    | ToolIdStateKey
+    | 'aws.amazonq.codewhisperer.newCustomizations'
+    | 'aws.amazonq.hasShownWalkthrough'
     | 'aws.downloadPath'
     | 'aws.lastTouchedS3Folder'
     | 'aws.lastUploadedToS3Folder'
     | 'aws.redshift.connections'
     | 'aws.toolkit.amazonq.dismissed'
     | 'aws.toolkit.amazonqInstall.dismissed'
-    | 'aws.toolkit.separationPromptDismissed'
     | 'aws.toolkit.separationPromptCommand'
-    | 'aws.amazonq.codewhisperer.newCustomizations'
+    | 'aws.toolkit.separationPromptDismissed'
     // Deprecated/legacy names. New keys should start with "aws.".
     | 'CODECATALYST_RECONNECT'
     | 'CODEWHISPERER_AUTO_SCANS_ENABLED'
@@ -38,6 +43,8 @@ type globalKey =
     | 'CODEWHISPERER_SELECTED_CUSTOMIZATION'
     | 'CODEWHISPERER_USER_GROUP'
     | 'gumby.wasQCodeTransformationUsed'
+    | 'isExtensionFirstUse'
+    | 'lastExtensionVersion'
     | 'hasAlreadyOpenedAmazonQ'
     // Legacy name from `ssoAccessTokenProvider.ts`.
     | '#sessionCreationDates'
diff --git a/packages/core/src/test/credentials/secondaryAuth.test.ts b/packages/core/src/test/credentials/secondaryAuth.test.ts
@@ -8,6 +8,7 @@ import { SecondaryAuth, getSecondaryAuth } from '../../auth/secondaryAuth'
 import { createBuilderIdProfile, createTestAuth } from './testUtil'
 import { Connection, createSsoProfile, hasScopes, isSsoConnection } from '../../auth/connection'
 import assert from 'assert'
+import globals from '../../shared/extensionGlobals'
 
 describe('SecondaryAuth', function () {
     let auth: ReturnType<typeof createTestAuth>
@@ -29,6 +30,7 @@ describe('SecondaryAuth', function () {
         secondaryAuth = getSecondaryAuth(auth, 'testId', 'testLabel', isValid)
         onDidChangeActiveConnection = sandbox.stub()
         secondaryAuth.onDidChangeActiveConnection(onDidChangeActiveConnection)
+        await globals.globalState.clear()
     })
 
     afterEach(async function () {
diff --git a/packages/core/src/test/credentials/utils.test.ts b/packages/core/src/test/credentials/utils.test.ts
@@ -2,60 +2,57 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import * as vscode from 'vscode'
 import assert from 'assert'
-import { FakeExtensionContext } from '../fakeExtensionContext'
 import { BuilderIdKind, ExtensionUse, SsoKind, hasBuilderId, hasIamCredentials, hasSso } from '../../auth/utils'
 import { Connection, SsoConnection, scopesCodeCatalyst } from '../../auth/connection'
 import { builderIdConnection, iamConnection, ssoConnection } from './testUtil'
 import { amazonQScopes } from '../../codewhisperer/util/authUtil'
+import globals from '../../shared/extensionGlobals'
 
 describe('ExtensionUse.isFirstUse()', function () {
-    let fakeState: vscode.Memento
     let instance: ExtensionUse
 
     beforeEach(async function () {
-        fakeState = (await FakeExtensionContext.create()).globalState
         instance = new ExtensionUse()
-        await fakeState.update(ExtensionUse.instance.isExtensionFirstUseKey, true)
+        await globals.globalState.update(ExtensionUse.instance.isExtensionFirstUseKey, true)
     })
 
     it('is true only on first startup', function () {
-        assert.strictEqual(instance.isFirstUse(fakeState), true, 'Failed on first call.')
-        assert.strictEqual(instance.isFirstUse(fakeState), true, 'Failed on second call.')
+        assert.strictEqual(instance.isFirstUse(), true, 'Failed on first call.')
+        assert.strictEqual(instance.isFirstUse(), true, 'Failed on second call.')
 
         const nextStartup = nextExtensionStartup()
-        assert.strictEqual(nextStartup.isFirstUse(fakeState), false, 'Failed on new startup.')
+        assert.strictEqual(nextStartup.isFirstUse(), false, 'Failed on new startup.')
     })
 
     it('true when: (state value not exists + NOT has existing connections)', async function () {
         await makeStateValueNotExist()
         const notHasExistingConnections = () => false
         assert.strictEqual(
-            instance.isFirstUse(fakeState, notHasExistingConnections),
+            instance.isFirstUse(notHasExistingConnections),
             true,
             'No existing connections, should be first use'
         )
-        assert.strictEqual(nextExtensionStartup().isFirstUse(fakeState), false)
+        assert.strictEqual(nextExtensionStartup().isFirstUse(), false)
     })
 
     it('false when: (state value not exists + has existing connections)', async function () {
         await makeStateValueNotExist()
         const hasExistingConnections = () => true
         assert.strictEqual(
-            instance.isFirstUse(fakeState, hasExistingConnections),
+            instance.isFirstUse(hasExistingConnections),
             false,
             'Found existing connections, should not be first use'
         )
-        assert.strictEqual(nextExtensionStartup().isFirstUse(fakeState), false)
+        assert.strictEqual(nextExtensionStartup().isFirstUse(), false)
     })
 
     /**
      * This makes the backend state value: undefined, mimicking a brand new user.
      * We use this state value to track if user is a first time user.
      */
     async function makeStateValueNotExist() {
-        await fakeState.update(ExtensionUse.instance.isExtensionFirstUseKey, undefined)
+        await globals.globalState.update(ExtensionUse.instance.isExtensionFirstUseKey, undefined)
     }
 
     /**

Original file line number	Diff line number	Diff line change
`@@ -115,7 +115,7 @@ describe('codewhisperer', async function () {`
`115`	`115`	`ideCategory: 'VSCODE',`
`116`	`116`	`operatingSystem: getOperatingSystem(),`
`117`	`117`	`product: 'CodeWhisperer',`
`118`		`- clientId: getClientId(globals.context.globalState),`
	`118`	`+ clientId: getClientId(globals.globalState),`
`119`	`119`	`}`
`120`	`120`
`121`	`121`	`await codeWhispererClient.sendTelemetryEvent({ telemetryEvent: userTriggerDecisionPayload })`
Original file line number	Diff line number	Diff line change
`@@ -337,7 +337,7 @@ export class FeatureDevClient {`
`337`	`337`	`ideCategory: 'VSCODE',`
`338`	`338`	`operatingSystem: getOperatingSystem(),`
`339`	`339`	`product: 'FeatureDev', // Should be the same as in JetBrains`
`340`		`- clientId: getClientId(globals.context.globalState),`
	`340`	`+ clientId: getClientId(globals.globalState),`
`341`	`341`	`ideVersion: extensionVersion,`
`342`	`342`	`},`
`343`	`343`	`}`