Add command to reset state (#3421)

* Add command to reset state

* Add tests

extension/package.json

@@ -558,6 +558,11 @@
         "command": "dvc.views.plotsPathsTree.refreshPlots",
         "category": "DVC",
         "icon": "$(refresh)"
+      },
+      {
+        "title": "Reset Persisted State and Reload Window",
+        "command": "dvc.resetState",
+        "category": "DVC"
       }
     ],
     "configuration": {

extension/src/commands/external.ts

@@ -101,5 +101,7 @@ export enum RegisteredCommands {
   ADD_STUDIO_ACCESS_TOKEN = 'dvc.addStudioAccessToken',
   UPDATE_STUDIO_ACCESS_TOKEN = 'dvc.updateStudioAccessToken',
   REMOVE_STUDIO_ACCESS_TOKEN = 'dvc.removeStudioAccessToken',
-  EXPERIMENT_VIEW_SHARE_TO_STUDIO = 'dvc.views.experiments.shareExperimentToStudio'
+  EXPERIMENT_VIEW_SHARE_TO_STUDIO = 'dvc.views.experiments.shareExperimentToStudio',
+
+  RESET_STATE = 'dvc.resetState'
 }

extension/src/extension.ts

@@ -56,6 +56,7 @@ import { LanguageClient } from './languageClient'
 import { collectRunningExperimentPids } from './experiments/processExecution/collect'
 import { registerPatchCommand } from './patch'
 import { DvcViewer } from './cli/dvc/viewer'
+import { registerPersistenceCommands } from './persistence/register'
 export class Extension extends Disposable {
   protected readonly internalCommands: InternalCommands
 
@@ -287,6 +288,8 @@ export class Extension extends Disposable {
       ).contributes.walkthroughs[0].id
     )
 
+    registerPersistenceCommands(context.workspaceState, this.internalCommands)
+
     void showWalkthroughOnFirstUse(env.isNewAppInstall)
     this.dispose.track(recommendRedHatExtensionOnce())
 

extension/src/persistence/register.ts

@@ -0,0 +1,13 @@
+import { Memento } from 'vscode'
+import { resetPersistedState } from './util'
+import { RegisteredCommands } from '../commands/external'
+import { InternalCommands } from '../commands/internal'
+
+export const registerPersistenceCommands = (
+  workspaceState: Memento,
+  internalCommands: InternalCommands
+) => {
+  internalCommands.registerExternalCommand(RegisteredCommands.RESET_STATE, () =>
+    resetPersistedState(workspaceState)
+  )
+}

extension/src/persistence/util.test.ts

@@ -0,0 +1,49 @@
+import { commands } from 'vscode'
+import { PersistenceKey } from './constants'
+import { resetPersistedState } from './util'
+import { buildMockMemento } from '../test/util'
+import {
+  DEFAULT_HEIGHT,
+  DEFAULT_SECTION_NB_ITEMS_PER_ROW
+} from '../plots/webview/contract'
+
+jest.mock('vscode')
+
+const mockedCommands = jest.mocked(commands)
+mockedCommands.executeCommand = jest.fn()
+
+describe('Persistence util', () => {
+  beforeEach(() => {
+    jest.resetAllMocks()
+  })
+
+  describe('resetPersistedState', () => {
+    it('should reload the window', async () => {
+      const workspaceState = buildMockMemento()
+
+      await resetPersistedState(workspaceState)
+
+      expect(mockedCommands.executeCommand).toHaveBeenCalledWith(
+        'workbench.action.reloadWindow'
+      )
+    })
+
+    it('should reset all values from all dvc roots', async () => {
+      const persistedState = {
+        [PersistenceKey.PLOT_HEIGHT + 'root1']: DEFAULT_HEIGHT,
+        [PersistenceKey.PLOT_NB_ITEMS_PER_ROW + 'root2']:
+          DEFAULT_SECTION_NB_ITEMS_PER_ROW
+      }
+      const workspaceState = buildMockMemento(persistedState)
+
+      await resetPersistedState(workspaceState)
+
+      expect(
+        workspaceState.get(PersistenceKey.PLOT_HEIGHT + 'root1')
+      ).toBeUndefined()
+      expect(
+        workspaceState.get(PersistenceKey.PLOT_NB_ITEMS_PER_ROW + 'root2')
+      ).toBeUndefined()
+    })
+  })
+})

extension/src/persistence/util.ts

@@ -0,0 +1,9 @@
+import { commands, Memento } from 'vscode'
+
+export const resetPersistedState = async (workspaceState: Memento) => {
+  for (const persistenceKey of workspaceState.keys()) {
+    await workspaceState.update(persistenceKey, undefined)
+  }
+
+  await commands.executeCommand('workbench.action.reloadWindow')
+}

extension/src/telemetry/constants.ts

@@ -291,4 +291,6 @@ export interface IEventNamePropertyMapping {
   [EventName.ADD_STUDIO_ACCESS_TOKEN]: undefined
   [EventName.UPDATE_STUDIO_ACCESS_TOKEN]: undefined
   [EventName.REMOVE_STUDIO_ACCESS_TOKEN]: undefined
+
+  [EventName.RESET_STATE]: undefined
 }

extension/src/test/util/index.ts

@@ -16,8 +16,8 @@ export const buildMockMemento = (
   ({
     get: (key: string, defaultValue: unknown) => values[key] || defaultValue,
     keys: () => Object.keys(values),
-    update: (key: string, value: unknown) =>
-      new Promise(() => {
-        values[key] = value
-      })
+    update: (key: string, value: unknown) => {
+      values[key] = value
+      void Promise.resolve()
+    }
   } as unknown as Memento)