refactor(codewhisperer): use globalState abstraction

aws · Jul 16, 2024 · afa70d1 · afa70d1
1 parent 177a649
commit afa70d1
Show file tree

Hide file tree

Showing 8 changed files with 43 additions and 21 deletions.
diff --git a/packages/core/src/codewhisperer/models/model.ts b/packages/core/src/codewhisperer/models/model.ts
@@ -92,7 +92,7 @@ export class CodeSuggestionsState {
     async toggleSuggestions() {
         const autoTriggerEnabled = this.isSuggestionsEnabled()
         const toSet: boolean = !autoTriggerEnabled
-        globals.globalState.tryUpdate('CODEWHISPERER_AUTO_TRIGGER_ENABLED', toSet)
+        await globals.globalState.tryUpdate('CODEWHISPERER_AUTO_TRIGGER_ENABLED', toSet)
         this.#onDidChangeState.fire(toSet)
         return toSet
     }
@@ -131,7 +131,7 @@ export class CodeScansState {
     async toggleScans() {
         const autoScansEnabled = this.isScansEnabled()
         const toSet: boolean = !autoScansEnabled
-        globals.globalState.tryUpdate('CODEWHISPERER_AUTO_SCANS_ENABLED', toSet)
+        await globals.globalState.tryUpdate('CODEWHISPERER_AUTO_SCANS_ENABLED', toSet)
         this.#onDidChangeState.fire(toSet)
         return toSet
     }

diff --git a/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts b/packages/core/src/codewhisperer/ui/codeWhispererNodes.ts
@@ -29,21 +29,21 @@ import { submitFeedback } from '../../feedback/vue/submitFeedback'
 import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands'
 import { isWeb } from '../../shared/extensionGlobals'
 
-export function createAutoSuggestions(pause: boolean): DataQuickPickItem<'autoSuggestions'> {
+export function createAutoSuggestions(running: boolean): DataQuickPickItem<'autoSuggestions'> {
     const labelResume = localize('AWS.codewhisperer.resumeCodeWhispererNode.label', 'Resume Auto-Suggestions')
     const iconResume = getIcon('vscode-debug-start')
     const labelPause = localize('AWS.codewhisperer.pauseCodeWhispererNode.label', 'Pause Auto-Suggestions')
     const iconPause = getIcon('vscode-debug-pause')
 
     return {
         data: 'autoSuggestions',
-        label: pause ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
-        description: pause ? 'Currently RUNNING' : 'Currently PAUSED',
+        label: running ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
+        description: running ? 'Currently RUNNING' : 'Currently PAUSED',
         onClick: () => toggleCodeSuggestions.execute(placeholder, cwQuickPickSource),
     } as DataQuickPickItem<'autoSuggestions'>
 }
 
-export function createAutoScans(pause: boolean): DataQuickPickItem<'autoScans'> {
+export function createAutoScans(running: boolean): DataQuickPickItem<'autoScans'> {
     const labelResume = localize('AWS.codewhisperer.resumeCodeWhispererNode.label', 'Resume Auto-Scans')
     const iconResume = getIcon('vscode-debug-alt')
     const labelPause = localize('AWS.codewhisperer.pauseCodeWhispererNode.label', 'Pause Auto-Scans')
@@ -52,8 +52,8 @@ export function createAutoScans(pause: boolean): DataQuickPickItem<'autoScans'>
 
     return {
         data: 'autoScans',
-        label: pause ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
-        description: monthlyQuotaExceeded ? 'Monthly quota exceeded' : pause ? 'RUNNING' : 'PAUSED',
+        label: running ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
+        description: monthlyQuotaExceeded ? 'Monthly quota exceeded' : running ? 'RUNNING' : 'PAUSED',
         onClick: () => toggleCodeScans.execute(placeholder, cwQuickPickSource),
     } as DataQuickPickItem<'autoScans'>
 }

diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts
@@ -134,7 +134,7 @@ export const setSelectedCustomization = async (customization: Customization) =>
     selectedCustomizationObj[AuthUtil.instance.conn.label] = customization
     getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.conn.label}`)
 
-    globals.globalState.tryUpdate('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomizationObj)
+    await globals.globalState.tryUpdate('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomizationObj)
     vsCodeState.isFreeTierLimitReached = false
     await Commands.tryExecute('aws.amazonq.refreshStatusBar')
 }
@@ -161,15 +161,15 @@ export const setPersistedCustomizations = async (customizations: Customization[]
         {}
     )
     persistedCustomizationsObj[AuthUtil.instance.conn.label] = customizations
-    globals.globalState.tryUpdate('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsObj)
+    await globals.globalState.tryUpdate('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsObj)
 }
 
 export const getNewCustomizationsAvailable = () => {
     return globals.globalState.tryGet('aws.amazonq.codewhisperer.newCustomizations', Number, 0)
 }
 
 export const setNewCustomizationsAvailable = async (num: number) => {
-    globals.globalState.tryUpdate('aws.amazonq.codewhisperer.newCustomizations', num)
+    await globals.globalState.tryUpdate('aws.amazonq.codewhisperer.newCustomizations', num)
     vsCodeState.isFreeTierLimitReached = false
 }
 

diff --git a/packages/core/src/codewhisperer/views/lineAnnotationController.ts b/packages/core/src/codewhisperer/views/lineAnnotationController.ts
@@ -308,7 +308,7 @@ export class LineAnnotationController implements vscode.Disposable {
     async dismissTutorial() {
         this._currentState = new EndState()
         await setContext('aws.codewhisperer.tutorial.workInProgress', false)
-        globals.globalState.tryUpdate(inlinehintKey, this._currentState.id)
+        await globals.globalState.tryUpdate(inlinehintKey, this._currentState.id)
     }
 
     private async onActiveLinesChanged(e: LinesChangeEvent) {
@@ -422,7 +422,7 @@ export class LineAnnotationController implements vscode.Disposable {
 
         decorationOptions.range = range
 
-        globals.globalState.tryUpdate(inlinehintKey, this._currentState.id)
+        await globals.globalState.tryUpdate(inlinehintKey, this._currentState.id)
         await setContext('aws.codewhisperer.tutorial.workInProgress', true)
         editor.setDecorations(this.cwLineHintDecoration, [decorationOptions])
     }

diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts
@@ -58,6 +58,10 @@ export class GlobalState implements vscode.Memento {
         return this.memento.keys()
     }
 
+    values() {
+        return this.memento.keys().map((k) => this.memento.get(k))
+    }
+
     /**
      * Gets the value for `key` if it satisfies the `type` specification, or fails.
      *
@@ -115,8 +119,8 @@ export class GlobalState implements vscode.Memento {
     }
 
     /** Asynchronously updates globalState, or logs an error on failure. */
-    tryUpdate(key: globalKey, value: any): void {
-        this.memento.update(key, value).then(
+    tryUpdate(key: globalKey, value: any) {
+        return this.memento.update(key, value).then(
             undefined, // TODO: log.debug() ?
             (e) => {
                 getLogger().error('GlobalState: failed to set "%s": %s', key, (e as Error).message)
@@ -128,6 +132,10 @@ export class GlobalState implements vscode.Memento {
         return this.memento.update(key, value)
     }
 
+    clear() {
+        return Promise.allSettled(this.memento.keys().map((k) => this.memento.update(k, undefined)))
+    }
+
     /**
      * Stores Redshift connection info for the specified warehouse ARN.
      *

diff --git a/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts b/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts
@@ -400,10 +400,10 @@ describe('CodeWhisperer-basicCommands', function () {
         it('shows expected quick pick items when connected', async function () {
             sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false)
             sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
-            sinon.stub(CodeScansState.instance, 'isScansEnabled').returns(false)
+            CodeScansState.instance.setScansEnabled(false)
             getTestWindow().onDidShowQuickPick((e) => {
                 e.assertContainsItems(
-                    createAutoSuggestions(false),
+                    createAutoSuggestions(true),
                     createOpenReferenceLog(),
                     createGettingStarted(),
                     createAutoScans(false),
@@ -423,11 +423,11 @@ describe('CodeWhisperer-basicCommands', function () {
             sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
             sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(true)
             sinon.stub(AuthUtil.instance, 'isCustomizationFeatureEnabled').value(true)
-            sinon.stub(CodeScansState.instance, 'isScansEnabled').returns(false)
+            CodeScansState.instance.setScansEnabled(false)
 
             getTestWindow().onDidShowQuickPick(async (e) => {
                 e.assertContainsItems(
-                    createAutoSuggestions(false),
+                    createAutoSuggestions(true),
                     createOpenReferenceLog(),
                     createGettingStarted(),
                     createAutoScans(false),
@@ -451,7 +451,7 @@ describe('CodeWhisperer-basicCommands', function () {
             getTestWindow().onDidShowQuickPick(async (e) => {
                 e.assertItems([
                     createSeparator('Inline Suggestions'),
-                    createAutoSuggestions(false),
+                    createAutoSuggestions(true),
                     createOpenReferenceLog(),
                     createGettingStarted(),
                     createSeparator('Security Scans'),

diff --git a/packages/core/src/test/codewhisperer/testUtil.ts b/packages/core/src/test/codewhisperer/testUtil.ts
@@ -32,7 +32,8 @@ export async function resetCodeWhispererGlobalVariables() {
     CodeWhispererCodeCoverageTracker.instances.clear()
     globals.telemetry.logger.clear()
     session.reset()
-    await CodeSuggestionsState.instance.setSuggestionsEnabled(false)
+    await globals.globalState.clear()
+    await CodeSuggestionsState.instance.setSuggestionsEnabled(true)
     await RecommendationHandler.instance.clearInlineCompletionStates()
 }
 

diff --git a/packages/core/src/test/shared/globalState.test.ts b/packages/core/src/test/shared/globalState.test.ts
@@ -97,6 +97,19 @@ describe('GlobalState', function () {
         assert.deepStrictEqual(globalState.tryGet(testKey, Boolean, true), true)
     })
 
+    it('clear()', async () => {
+        const keys = ['CODECATALYST_RECONNECT', 'SAM_INIT_ARCH_KEY', 'aws.redshift.connections']
+        await globalState.update(keys[0] as any, 'val1')
+        await globalState.update(keys[1] as any, 'val2')
+        await globalState.update(keys[2] as any, 'val3')
+        assert.deepStrictEqual(globalState.keys(), keys)
+        assert.deepStrictEqual(globalState.values(), ['val1', 'val2', 'val3'])
+        await globalState.clear()
+        // XXX: no way to actually delete the key?
+        assert.deepStrictEqual(globalState.keys(), keys)
+        assert.deepStrictEqual(globalState.values(), [undefined, undefined, undefined])
+    })
+
     describe('redshift state', function () {
         const testArn1 = 'arn:foo/bar/baz/1'
         const testArn2 = 'arn:foo/bar/baz/2'