Tests for undo using keyboard shortcut 'z' (microsoft#7824)

DonJayamanne · web-flow · commit cfb0d401af8b · 2019-10-08T11:13:10.000-07:00
diff --git a/src/datascience-ui/interactive-common/mainStateController.ts b/src/datascience-ui/interactive-common/mainStateController.ts
@@ -517,11 +517,13 @@ export class MainStateController implements IMessageHandler {
         // Skip if already has focus
         if (cellId !== this.pendingState.focusedCellId) {
             const newVMs = [...this.pendingState.cellVMs];
-            // Switch the old vm
-            const oldSelect = this.findCellIndex(this.pendingState.selectedCellId);
-            if (oldSelect >= 0) {
-                newVMs[oldSelect] = { ...newVMs[oldSelect], selected: false, focused: false };
-            }
+            // Reset the old vms (nothing should be selected/focused)
+            // Change state only for cells that were selected/focused
+            newVMs.forEach((cellVM, index) => {
+                if (cellVM.selected || cellVM.focused){
+                    newVMs[index] = { ...cellVM, selected: false, focused: false };
+                }
+            });
             const newSelect = this.findCellIndex(cellId);
             if (newSelect >= 0) {
                 newVMs[newSelect] = { ...newVMs[newSelect], selected: true, focused: true };
@@ -536,11 +538,13 @@ export class MainStateController implements IMessageHandler {
         // Skip if already the same cell
         if (this.pendingState.selectedCellId !== cellId || this.pendingState.focusedCellId !== focusedCellId) {
             const newVMs = [...this.pendingState.cellVMs];
-            // Switch the old vm
-            const oldSelect = this.findCellIndex(this.pendingState.selectedCellId);
-            if (oldSelect >= 0) {
-                newVMs[oldSelect] = { ...newVMs[oldSelect], selected: false, focused: false };
-            }
+            // Reset the old vms (nothing should be selected/focused)
+            // Change state only for cells that were selected/focused
+            newVMs.forEach((cellVM, index) => {
+                if (cellVM.selected || cellVM.focused){
+                    newVMs[index] = { ...cellVM, selected: false, focused: false };
+                }
+            });
             const newSelect = this.findCellIndex(cellId);
             if (newSelect >= 0) {
                 newVMs[newSelect] = { ...newVMs[newSelect], selected: true, focused: focusedCellId === newVMs[newSelect].cell.id };
diff --git a/src/test/datascience/nativeEditor.functional.test.tsx b/src/test/datascience/nativeEditor.functional.test.tsx
@@ -659,5 +659,34 @@ suite('DataScience Native Editor', () => {
             assert.equal(wrapper.find(NativeCell).at(1).find(CellOutput).length, 0);
             assert.equal(wrapper.find(NativeCell).at(1).find(MonacoEditor).length, 1);
         });
+
+        test('Test undo using the key \'z\'', async () => {
+            clickCell(0);
+
+            // Add, then undo, keep doing at least 3 times and confirm it works as expected.
+            for (let i = 0; i < 3; i += 1){
+                // Add a new cell
+                let update = waitForUpdate(wrapper, NativeEditor, 1);
+                simulateKeyPressOnCell(0, { code: 'a' });
+                await update;
+
+                // There should be 4 cells and first cell is selected & nothing focused.
+                assert.equal(isCellSelected(wrapper, 'NativeCell', 0), true);
+                assert.equal(isCellSelected(wrapper, 'NativeCell', 1), false);
+                assert.equal(isCellFocused(wrapper, 'NativeCell', 0), false);
+                assert.equal(isCellFocused(wrapper, 'NativeCell', 1), false);
+                assert.equal(wrapper.find('NativeCell').length, 4);
+
+                // Press 'z' to undo.
+                update = waitForUpdate(wrapper, NativeEditor, 1);
+                simulateKeyPressOnCell(0, { code: 'z' });
+                await update;
+
+                // There should be 3 cells and first cell is selected & nothing focused.
+                assert.equal(isCellSelected(wrapper, 'NativeCell', 0), true);
+                assert.equal(isCellSelected(wrapper, 'NativeCell', 1), false);
+                assert.equal(wrapper.find('NativeCell').length, 3);
+            }
+        });
     });
 });
