refactor(stepFunctions): use globalState abstraction

packages/core/src/shared/globalState.ts

@@ -15,8 +15,11 @@ type samInitStateKey =
     | 'SAM_INIT_IMAGE_BOOLEAN_KEY'
     | 'SAM_INIT_ARCH_KEY'
 
+type stepFunctionsKey = 'SCRIPT_LAST_DOWNLOADED_URL' | 'CSS_LAST_DOWNLOADED_URL'
+
 type globalKey =
     | samInitStateKey
+    | stepFunctionsKey
     | 'aws.downloadPath'
     | 'aws.lastTouchedS3Folder'
     | 'aws.lastUploadedToS3Folder'

packages/core/src/stepFunctions/activation.ts

@@ -62,7 +62,7 @@ export async function activate(
  */
 export const previewStateMachineCommand = Commands.declare(
     'aws.previewStateMachine',
-    (globalState: vscode.Memento, manager: AslVisualizationManager) => async (arg?: vscode.TextEditor | vscode.Uri) => {
+    (manager: AslVisualizationManager) => async (arg?: vscode.TextEditor | vscode.Uri) => {
         try {
             arg ??= vscode.window.activeTextEditor
             const input = arg instanceof vscode.Uri ? arg : arg?.document
@@ -71,7 +71,7 @@ export const previewStateMachineCommand = Commands.declare(
                 throw new ToolkitError('No active text editor or document found')
             }
 
-            return await manager.visualizeStateMachine(globalState, input)
+            return await manager.visualizeStateMachine(input)
         } finally {
             // TODO: Consider making the metric reflect the success/failure of the above call
             telemetry.stepfunctions_previewstatemachine.emit()
@@ -88,8 +88,8 @@ async function registerStepFunctionCommands(
     const cdkVisualizationManager = new AslVisualizationCDKManager(extensionContext)
 
     extensionContext.subscriptions.push(
-        previewStateMachineCommand.register(extensionContext.globalState, visualizationManager),
-        renderCdkStateMachineGraph.register(extensionContext.globalState, cdkVisualizationManager),
+        previewStateMachineCommand.register(visualizationManager),
+        renderCdkStateMachineGraph.register(cdkVisualizationManager),
         Commands.register('aws.stepfunctions.createStateMachineFromTemplate', async () => {
             try {
                 await createStateMachineFromTemplate(extensionContext)

packages/core/src/stepFunctions/commands/visualizeStateMachine/abstractAslVisualizationManager.ts

@@ -19,10 +19,7 @@ export abstract class AbstractAslVisualizationManager<T extends AslVisualization
 
     public constructor(private readonly extensionContext: vscode.ExtensionContext) {}
 
-    public abstract visualizeStateMachine(
-        globalState: vscode.Memento,
-        uri: vscode.Uri
-    ): Promise<vscode.WebviewPanel | undefined>
+    public abstract visualizeStateMachine(uri: vscode.Uri): Promise<vscode.WebviewPanel | undefined>
 
     protected pushToExtensionContextSubscriptions(visualizationDisposable: vscode.Disposable): void {
         this.extensionContext.subscriptions.push(visualizationDisposable)
@@ -57,9 +54,9 @@ export abstract class AbstractAslVisualizationManager<T extends AslVisualization
         return this.managedVisualizations.get(key)
     }
 
-    protected async updateCache(globalState: vscode.Memento, logger: Logger): Promise<void> {
+    protected async updateCache(logger: Logger): Promise<void> {
         try {
-            await this.cache.updateCache(globalState)
+            await this.cache.updateCache()
         } catch (err) {
             // So we can't update the cache, but can we use an existing on disk version.
             logger.warn('Updating State Machine Graph Visualisation assets failed, checking for fallback local cache.')

packages/core/src/stepFunctions/commands/visualizeStateMachine/aslVisualizationCDKManager.ts

@@ -16,10 +16,7 @@ export class AslVisualizationCDKManager extends AbstractAslVisualizationManager<
         super(extensionContext)
     }
 
-    public async visualizeStateMachine(
-        globalState: vscode.Memento,
-        uri: vscode.Uri
-    ): Promise<vscode.WebviewPanel | undefined> {
+    public async visualizeStateMachine(uri: vscode.Uri): Promise<vscode.WebviewPanel | undefined> {
         const logger = getLogger()
         const existingVisualization = this.getExistingVisualization(this.getKey(uri))
 
@@ -34,7 +31,7 @@ export class AslVisualizationCDKManager extends AbstractAslVisualizationManager<
         const templateUri = vscode.Uri.joinPath(cdkOutPath, `${appName}.template.json`)
 
         try {
-            await this.cache.updateCache(globalState)
+            await this.cache.updateCache()
 
             const textDocument = await vscode.workspace.openTextDocument(templateUri.with({ fragment: '' }))
             const newVisualization = new AslVisualizationCDK(textDocument, templateUri.fsPath, resourceName)

packages/core/src/stepFunctions/commands/visualizeStateMachine/aslVisualizationManager.ts

@@ -17,7 +17,6 @@ export class AslVisualizationManager extends AbstractAslVisualizationManager {
     }
 
     public async visualizeStateMachine(
-        globalState: vscode.Memento,
         target: vscode.TextDocument | vscode.Uri
     ): Promise<vscode.WebviewPanel | undefined> {
         const logger: Logger = getLogger()
@@ -33,7 +32,7 @@ export class AslVisualizationManager extends AbstractAslVisualizationManager {
 
         // Existing visualization does not exist, construct new visualization
         try {
-            await this.updateCache(globalState, logger)
+            await this.updateCache(logger)
             const newVisualization = new AslVisualization(document)
             this.handleNewVisualization(document.uri.fsPath, newVisualization)
 

packages/core/src/stepFunctions/commands/visualizeStateMachine/renderStateMachineGraphCDK.ts

@@ -23,14 +23,14 @@ function isLocationResource(obj: unknown): obj is { location: vscode.Uri } {
  */
 export const renderCdkStateMachineGraph = Commands.declare(
     'aws.cdk.renderStateMachineGraph',
-    (memento: vscode.Memento, manager: AslVisualizationCDKManager) => async (input?: unknown) => {
+    (manager: AslVisualizationCDKManager) => async (input?: unknown) => {
         const resource = isTreeNode(input) ? unboxTreeNode(input, isLocationResource) : undefined
         const resourceUri = resource?.location ?? (await new PreviewStateMachineCDKWizard().run())?.resource.location
 
         if (!resourceUri) {
             return
         }
 
-        await manager.visualizeStateMachine(memento, resourceUri)
+        await manager.visualizeStateMachine(resourceUri)
     }
 )

packages/core/src/stepFunctions/utils.ts

@@ -31,8 +31,7 @@ const scriptsLastDownloadedUrl = 'SCRIPT_LAST_DOWNLOADED_URL'
 const cssLastDownloadedUrl = 'CSS_LAST_DOWNLOADED_URL'
 
 export interface UpdateCachedScriptOptions {
-    globalState: vscode.Memento
-    lastDownloadedURLKey: string
+    lastDownloadedURLKey: 'SCRIPT_LAST_DOWNLOADED_URL' | 'CSS_LAST_DOWNLOADED_URL'
     currentURL: string
     filePath: string
 }
@@ -73,9 +72,8 @@ export class StateMachineGraphCache {
         this.fileExists = fileExistsCustom ?? fileExists
     }
 
-    public async updateCache(globalState: vscode.Memento): Promise<void> {
+    public async updateCache(): Promise<void> {
         const scriptUpdate = this.updateCachedFile({
-            globalState: globalState,
             lastDownloadedURLKey: scriptsLastDownloadedUrl,
             currentURL: visualizationScriptUrl,
             filePath: this.jsFilePath,
@@ -87,7 +85,6 @@ export class StateMachineGraphCache {
         })
 
         const cssUpdate = this.updateCachedFile({
-            globalState: globalState,
             lastDownloadedURLKey: cssLastDownloadedUrl,
             currentURL: visualizationCssUrl,
             filePath: this.cssFilePath,
@@ -102,7 +99,7 @@ export class StateMachineGraphCache {
     }
 
     public async updateCachedFile(options: UpdateCachedScriptOptions) {
-        const downloadedUrl = options.globalState.get<string>(options.lastDownloadedURLKey)
+        const downloadedUrl = globals.globalState.tryGet<string>(options.lastDownloadedURLKey, String)
         const cachedFileExists = await this.fileExists(options.filePath)
 
         // if current url is different than url that was previously used to download the assets
@@ -113,7 +110,7 @@ export class StateMachineGraphCache {
             await this.writeToLocalStorage(options.filePath, response)
 
             // save the url of the downloaded and cached assets
-            void options.globalState.update(options.lastDownloadedURLKey, options.currentURL)
+            await globals.globalState.update(options.lastDownloadedURLKey, options.currentURL)
         }
     }
 

packages/core/src/test/stepFunctions/commands/visualizeStateMachine.test.ts

@@ -90,7 +90,7 @@ describe('StepFunctions VisualizeStateMachine', async function () {
         const doc = await getDoc1()
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc)
+        await aslVisualizationManager.visualizeStateMachine(doc)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 1)
 
         const managedVisualizations = aslVisualizationManager.getManagedVisualizations()
@@ -101,10 +101,10 @@ describe('StepFunctions VisualizeStateMachine', async function () {
         const doc = await getDoc1()
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc)
+        await aslVisualizationManager.visualizeStateMachine(doc)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 1)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc)
+        await aslVisualizationManager.visualizeStateMachine(doc)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 1)
 
         const managedVisualizations = aslVisualizationManager.getManagedVisualizations()
@@ -117,10 +117,10 @@ describe('StepFunctions VisualizeStateMachine', async function () {
 
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc1)
+        await aslVisualizationManager.visualizeStateMachine(doc1)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 1)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc2)
+        await aslVisualizationManager.visualizeStateMachine(doc2)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 2)
 
         const managedVisualizations = aslVisualizationManager.getManagedVisualizations()
@@ -134,16 +134,16 @@ describe('StepFunctions VisualizeStateMachine', async function () {
 
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc1)
+        await aslVisualizationManager.visualizeStateMachine(doc1)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 1)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc2)
+        await aslVisualizationManager.visualizeStateMachine(doc2)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 2)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc1)
+        await aslVisualizationManager.visualizeStateMachine(doc1)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 2)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc2)
+        await aslVisualizationManager.visualizeStateMachine(doc2)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 2)
 
         const managedVisualizations = aslVisualizationManager.getManagedVisualizations()
@@ -154,7 +154,7 @@ describe('StepFunctions VisualizeStateMachine', async function () {
     it('Test AslVisualizationManager managedVisualizations set removes visualization on visualization dispose, single vis', async function () {
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
-        let panel = await aslVisualizationManager.visualizeStateMachine(globals.globalState, await getDoc1())
+        let panel = await aslVisualizationManager.visualizeStateMachine(await getDoc1())
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 1)
 
         // Dispose of visualization panel
@@ -171,10 +171,10 @@ describe('StepFunctions VisualizeStateMachine', async function () {
 
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
-        let panel = await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc1)
+        let panel = await aslVisualizationManager.visualizeStateMachine(doc1)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 1)
 
-        await aslVisualizationManager.visualizeStateMachine(globals.globalState, doc2)
+        await aslVisualizationManager.visualizeStateMachine(doc2)
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 2)
 
         // Dispose of first visualization panel

packages/core/src/test/stepFunctions/utils.test.ts

@@ -14,7 +14,7 @@ import { isDocumentValid, isStepFunctionsRole, StateMachineGraphCache } from '..
 const requestBody = 'request body string'
 const assetUrl = 'https://something'
 const filePath = '/some/path'
-const storageKey = 'KEY'
+const storageKey = 'SCRIPT_LAST_DOWNLOADED_URL'
 let tempFolder = ''
 
 describe('StateMachineGraphCache', function () {
@@ -49,7 +49,6 @@ describe('StateMachineGraphCache', function () {
             })
 
             await cache.updateCachedFile({
-                globalState: globalStorage,
                 lastDownloadedURLKey: storageKey,
                 currentURL: assetUrl,
                 filePath: filePath,
@@ -81,7 +80,6 @@ describe('StateMachineGraphCache', function () {
             })
 
             await cache.updateCachedFile({
-                globalState: globalStorage,
                 lastDownloadedURLKey: storageKey,
                 currentURL: assetUrl,
                 filePath: filePath,
@@ -113,7 +111,6 @@ describe('StateMachineGraphCache', function () {
             })
 
             await cache.updateCachedFile({
-                globalState: globalStorage,
                 lastDownloadedURLKey: storageKey,
                 currentURL: assetUrl,
                 filePath: filePath,
@@ -185,7 +182,6 @@ describe('StateMachineGraphCache', function () {
             })
 
             await cache.updateCachedFile({
-                globalState: globalStorage,
                 lastDownloadedURLKey: storageKey,
                 currentURL: assetUrl,
                 filePath: filePath,