feat(context): add getContext, update notifications tests

Add counterpart to setContext wrapper.

Author: hayemaxi

packages/core/src/notifications/controller.ts

@@ -92,7 +92,7 @@ export class NotificationsController {
             ruleEngine.shouldDisplayNotification(n)
         )
 
-        NotificationsNode.instance.setNotifications(startUp, emergency)
+        await NotificationsNode.instance.setNotifications(startUp, emergency)
 
         // Emergency notifications can't be dismissed, but if the user minimizes the panel then
         // we don't want to focus it each time we set the notification nodes.
@@ -126,7 +126,7 @@ export class NotificationsController {
         this.state.dismissed.push(notificationId)
         await this.writeState()
 
-        NotificationsNode.instance.dismissStartUpNotification(notificationId)
+        await NotificationsNode.instance.dismissStartUpNotification(notificationId)
     }
 
     /**

packages/core/src/notifications/panelNode.ts

@@ -77,7 +77,7 @@ export class NotificationsNode implements TreeNode {
         return item
     }
 
-    public refresh(): void {
+    public refresh() {
         const totalNotifications = this.notificationCount()
         if (this.view) {
             if (totalNotifications > 0) {
@@ -94,8 +94,8 @@ export class NotificationsNode implements TreeNode {
             logger.warn('NotificationsNode was refreshed but the view was not initialized!')
         }
 
-        void setContext(this.showContextStr, totalNotifications > 0)
         this.provider?.refresh()
+        return setContext(this.showContextStr, totalNotifications > 0)
     }
 
     public getChildren() {
@@ -126,10 +126,10 @@ export class NotificationsNode implements TreeNode {
      * Sets the current list of notifications. Nodes are generated for each notification.
      * No other processing is done, see NotificationController.
      */
-    public setNotifications(startUp: ToolkitNotification[], emergency: ToolkitNotification[]) {
+    public async setNotifications(startUp: ToolkitNotification[], emergency: ToolkitNotification[]) {
         this.startUpNotifications = startUp
         this.emergencyNotifications = emergency
-        this.refresh()
+        await this.refresh()
     }
 
     /**
@@ -138,9 +138,9 @@ export class NotificationsNode implements TreeNode {
      *
      * Only dismisses startup notifications.
      */
-    public dismissStartUpNotification(id: string) {
+    public async dismissStartUpNotification(id: string) {
         this.startUpNotifications = this.startUpNotifications.filter((n) => n.id !== id)
-        this.refresh()
+        await this.refresh()
     }
 
     /**

packages/core/src/shared/vscode/setContext.ts

@@ -36,22 +36,34 @@ export type contextKey =
     | 'amazonq.inline.codelensShortcutEnabled'
     | 'aws.toolkit.lambda.walkthroughSelected'
 
+const contextMap: Partial<Record<contextKey, any>> = {}
+
 /**
  * Calls the vscode "setContext" command.
  *
  * This wrapper adds structure and traceability to the vscode "setContext". It also opens the door
- * for:
- * - validation
- * - getContext() (see also https://github.com/microsoft/vscode/issues/10471)
+ * for validation.
  *
  * Use "setContext" only as a last resort, to set flags that are detectable in package.json
  * declarations. Do not use it as a general way to store global state (which should be avoided
  * anyway).
  *
  * Warning: vscode context keys/values are NOT isolated to individual extensions. Other extensions
- * can read and modify them.
+ * can read and modify them. See also https://github.com/microsoft/vscode/issues/10471
  */
 export async function setContext(key: contextKey, val: any): Promise<void> {
     // eslint-disable-next-line aws-toolkits/no-banned-usages
     await vscode.commands.executeCommand('setContext', key, val)
+    contextMap[key] = val
+}
+
+/**
+ * Returns the value of a context key set via {@link setContext} wrapper for this session.
+ *
+ * Warning: this does not gaurantee the state of the context key in vscode because it may have
+ * been set via `vscode.commands.executeCommand('setContext')`. It has no connection the
+ * context keys stored in vscode itself because an API for this is not exposed.
+ */
+export function getContext(key: contextKey): any {
+    return contextMap[key]
 }

packages/core/src/test/notifications/controller.test.ts

@@ -9,7 +9,7 @@ import assert from 'assert'
 import sinon from 'sinon'
 import globals from '../../shared/extensionGlobals'
 import { randomUUID } from '../../shared/crypto'
-import * as setContext from '../../shared/vscode/setContext'
+import { getContext } from '../../shared/vscode/setContext'
 import { assertTelemetry, installFakeClock } from '../testUtil'
 import {
     NotificationFetcher,
@@ -88,7 +88,7 @@ describe('Notifications Controller', function () {
     }
 
     beforeEach(async function () {
-        panelNode.setNotifications([], [])
+        await panelNode.setNotifications([], [])
         fetcher = new TestFetcher()
         controller = new NotificationsController(panelNode, fetcher, '_aws.test.notification' as any)
 
@@ -108,10 +108,6 @@ describe('Notifications Controller', function () {
     })
 
     it('can fetch and store startup notifications', async function () {
-        // There seems to be a race condition with having a global spy object that is reset after each
-        // test. Doesn't seem to affect the other ones, for some reason.
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -140,14 +136,10 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.startUpNotifications, [content.notifications[0]])
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(focusPanelSpy.callCount, 0)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('can fetch and store emergency notifications', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -176,14 +168,10 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.emergencyNotifications, [content.notifications[0]])
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(focusPanelSpy.callCount, 1)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('can fetch and store both startup and emergency notifications', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag1 = randomUUID()
         const eTag2 = randomUUID()
         const startUpContent = {
@@ -239,14 +227,10 @@ describe('Notifications Controller', function () {
         assert.deepStrictEqual(panelNode.emergencyNotifications, [emergencyContent.notifications[0]])
         assert.equal(panelNode.getChildren().length, 2)
         assert.equal(focusPanelSpy.callCount, 1)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('dismisses a startup notification', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const eTag = randomUUID()
         const content = {
             schemaVersion: '1.x',
@@ -261,7 +245,7 @@ describe('Notifications Controller', function () {
 
         assert.equal(panelNode.getChildren().length, 2)
         assert.equal(panelNode.startUpNotifications.length, 2)
-        assert.ok(setContextSpy.calledWithExactly('aws.toolkit.notifications.show', true))
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
 
         assert.deepStrictEqual(await globals.globalState.get(controller.storageKey), {
             startUp: {
@@ -288,13 +272,9 @@ describe('Notifications Controller', function () {
 
         assert.equal(panelNode.getChildren().length, 1)
         assert.equal(panelNode.startUpNotifications.length, 1)
-
-        setContextSpy.restore()
     })
 
     it('does not redisplay dismissed notifications', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const content = {
             schemaVersion: '1.x',
             notifications: [getValidTestNotification('id:startup1')],
@@ -306,11 +286,11 @@ describe('Notifications Controller', function () {
 
         await controller.pollForStartUp(ruleEngine)
         assert.equal(panelNode.getChildren().length, 1)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
 
         await dismissNotification(content.notifications[0])
         assert.equal(panelNode.getChildren().length, 0)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', false])
+        assert.equal(getContext('aws.toolkit.notifications.show'), false)
 
         content.notifications.push(getValidTestNotification('id:startup2'))
         fetcher.setStartUpContent({
@@ -332,9 +312,7 @@ describe('Notifications Controller', function () {
         })
 
         assert.equal(panelNode.getChildren().length, 1)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
     })
 
     it('does not refocus emergency notifications', async function () {
@@ -406,8 +384,6 @@ describe('Notifications Controller', function () {
     })
 
     it('cleans out dismissed state', async function () {
-        const setContextSpy = sinon.spy(setContext, 'setContext')
-
         const startUpContent = {
             schemaVersion: '1.x',
             notifications: [getValidTestNotification('id:startup1')],
@@ -466,7 +442,7 @@ describe('Notifications Controller', function () {
             newlyReceived: [],
         })
         assert.equal(panelNode.getChildren().length, 1)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', true])
+        assert.equal(getContext('aws.toolkit.notifications.show'), true)
 
         fetcher.setEmergencyContent({
             eTag: '1',
@@ -488,9 +464,7 @@ describe('Notifications Controller', function () {
         })
 
         assert.equal(panelNode.getChildren().length, 0)
-        assert.deepStrictEqual(setContextSpy.lastCall.args, ['aws.toolkit.notifications.show', false])
-
-        setContextSpy.restore()
+        assert.equal(getContext('aws.toolkit.notifications.show'), false)
     })
 
     it('does not rethrow errors when fetching', async function () {