Fix crash induced by webview disposal within tests (#1098)

Visualization panel tests caused a bug causing an extension host crash, when a webview was disposed two times. This change updates the code flow to prevent this.

Author: jormello

src/stepFunctions/commands/visualizeStateMachine/aslVisualization.ts

@@ -86,7 +86,7 @@ export class AslVisualization {
         // If documentUri being tracked is no longer found (due to file closure or rename), close the panel.
         this.disposables.push(
             vscode.workspace.onDidCloseTextDocument(documentWillSaveEvent => {
-                if (!this.trackedDocumentDoesExist(documentUri)) {
+                if (!this.trackedDocumentDoesExist(documentUri) && !this.isPanelDisposed) {
                     panel.dispose()
                     vscode.window.showInformationMessage(
                         localize(
@@ -153,15 +153,24 @@ export class AslVisualization {
         )
 
         // When the panel is closed, dispose of any disposables/remove subscriptions
-        panel.onDidDispose(() => {
+        const disposePanel = () => {
+            if (this.isPanelDisposed) {
+                return
+            }
+            this.isPanelDisposed = true
             debouncedUpdate.cancel()
             this.onVisualizationDisposeEmitter.fire()
             this.disposables.forEach(disposable => {
                 disposable.dispose()
             })
             this.onVisualizationDisposeEmitter.dispose()
-            this.isPanelDisposed = true
-        })
+        }
+
+        this.disposables.push(
+            panel.onDidDispose(() => {
+                disposePanel()
+            })
+        )
 
         return panel
     }

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

@@ -199,7 +199,7 @@ describe('StepFunctions VisualizeStateMachine', async () => {
 
         assert.deepStrictEqual(vis.documentUri, mockTextDocumentOne.uri)
         assert.strictEqual(vis.getIsPanelDisposed(), false)
-        assert.strictEqual(vis.getDisposables().length, 4)
+        assert.strictEqual(vis.getDisposables().length, 5)
 
         let panel = vis.getPanel()
         assert.ok(panel)
@@ -303,10 +303,7 @@ describe('StepFunctions VisualizeStateMachine', async () => {
         assert.ok(managedVisualizations.get(mockTextDocumentTwo.uri.path))
     })
 
-    // TODO re-enable these tests once this race condition is resolved.
-    // https://github.com/aws/aws-toolkit-vscode/issues/1071
-    //
-    it.skip('Test AslVisualizationManager managedVisualizations set removes visualization on visualization dispose, single vis', async () => {
+    it('Test AslVisualizationManager managedVisualizations set removes visualization on visualization dispose, single vis', async () => {
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
         // Preview Doc1
@@ -322,7 +319,7 @@ describe('StepFunctions VisualizeStateMachine', async () => {
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
     })
 
-    it.skip('Test AslVisualizationManager managedVisualizations set removes correct visualization on visualization dispose, multiple vis', async () => {
+    it('Test AslVisualizationManager managedVisualizations set removes correct visualization on visualization dispose, multiple vis', async () => {
         assert.strictEqual(aslVisualizationManager.getManagedVisualizations().size, 0)
 
         // Preview Doc1