Fix "show toast on new install" logic (#4024)

extension/src/extension.ts

@@ -1,4 +1,4 @@
-import { commands, env, ExtensionContext, ViewColumn } from 'vscode'
+import { commands, ExtensionContext, ViewColumn } from 'vscode'
 import { DvcConfig } from './cli/dvc/config'
 import { DvcExecutor } from './cli/dvc/executor'
 import { DvcRunner } from './cli/dvc/runner'
@@ -254,9 +254,13 @@ class Extension extends Disposable {
       ).contributes.walkthroughs[0].id
     )
 
-    registerPersistenceCommands(context.workspaceState, this.internalCommands)
+    registerPersistenceCommands(
+      context.workspaceState,
+      context.globalState,
+      this.internalCommands
+    )
 
-    void showSetupOnFirstUse(env.isNewAppInstall)
+    void showSetupOnFirstUse(context.globalState, context.workspaceState)
     this.dispose.track(recommendRedHatExtensionOnce())
 
     this.dispose.track(new LanguageClient())

extension/src/persistence/constants.ts

@@ -19,3 +19,7 @@ export enum PersistenceKey {
   PLOT_SELECTED_METRICS = 'plotSelectedMetrics:',
   PLOT_TEMPLATE_ORDER = 'plotTemplateOrder:'
 }
+
+export enum GlobalPersistenceKey {
+  INSTALLED = 'dvc.installed'
+}

extension/src/persistence/register.ts

@@ -5,9 +5,10 @@ import { InternalCommands } from '../commands/internal'
 
 export const registerPersistenceCommands = (
   workspaceState: Memento,
+  globalState: Memento,
   internalCommands: InternalCommands
 ) => {
   internalCommands.registerExternalCommand(RegisteredCommands.RESET_STATE, () =>
-    resetPersistedState(workspaceState)
+    resetPersistedState(workspaceState, globalState)
   )
 }

extension/src/persistence/util.test.ts

@@ -1,5 +1,5 @@
 import { commands } from 'vscode'
-import { PersistenceKey } from './constants'
+import { GlobalPersistenceKey, PersistenceKey } from './constants'
 import { resetPersistedState } from './util'
 import { buildMockMemento } from '../test/util'
 import {
@@ -20,23 +20,28 @@ describe('Persistence util', () => {
   describe('resetPersistedState', () => {
     it('should reload the window', async () => {
       const workspaceState = buildMockMemento()
+      const globalState = buildMockMemento()
 
-      await resetPersistedState(workspaceState)
+      await resetPersistedState(workspaceState, globalState)
 
       expect(mockedCommands.executeCommand).toHaveBeenCalledWith(
         'workbench.action.reloadWindow'
       )
     })
 
     it('should reset all values from all dvc roots', async () => {
-      const persistedState = {
+      const persistedWorkspaceState = {
         [PersistenceKey.PLOT_HEIGHT + 'root1']: DEFAULT_HEIGHT,
         [PersistenceKey.PLOT_NB_ITEMS_PER_ROW_OR_WIDTH + 'root2']:
           DEFAULT_SECTION_NB_ITEMS_PER_ROW_OR_WIDTH
       }
-      const workspaceState = buildMockMemento(persistedState)
+      const persistedGlobalState = {
+        [GlobalPersistenceKey.INSTALLED]: true
+      }
+      const workspaceState = buildMockMemento(persistedWorkspaceState)
+      const globalState = buildMockMemento(persistedGlobalState)
 
-      await resetPersistedState(workspaceState)
+      await resetPersistedState(workspaceState, globalState)
 
       expect(
         workspaceState.get(PersistenceKey.PLOT_HEIGHT + 'root1')
@@ -46,6 +51,7 @@ describe('Persistence util', () => {
           PersistenceKey.PLOT_NB_ITEMS_PER_ROW_OR_WIDTH + 'root2'
         )
       ).toBeUndefined()
+      expect(globalState.get(GlobalPersistenceKey.INSTALLED)).toBeUndefined()
     })
   })
 })

extension/src/persistence/util.ts

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

extension/src/setup/util.test.ts

@@ -4,6 +4,9 @@ import { ConfigKey, getConfigValue, setUserConfigValue } from '../vscode/config'
 import { Toast } from '../vscode/toast'
 import { Response } from '../vscode/response'
 import { RegisteredCommands } from '../commands/external'
+import { buildMockMemento } from '../test/util'
+import { GlobalPersistenceKey, PersistenceKey } from '../persistence/constants'
+import { DEFAULT_HEIGHT } from '../plots/webview/contract'
 
 jest.mock('vscode')
 jest.mock('../vscode/toast')
@@ -25,19 +28,51 @@ beforeEach(() => {
 })
 
 describe('showSetupOnFirstUse', () => {
+  it('should set a global state key after a new install', async () => {
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
+
+    expect(globalState.get(GlobalPersistenceKey.INSTALLED)).toStrictEqual(true)
+  })
+
   it('should ask to show the setup page after a new install', async () => {
-    await showSetupOnFirstUse(true)
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
+
     expect(mockedAskShowOrCloseOrNever).toHaveBeenCalledTimes(1)
   })
 
   it('should not ask to show the setup page when the install is not new', async () => {
-    await showSetupOnFirstUse(false)
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento({
+      [GlobalPersistenceKey.INSTALLED]: true
+    })
+
+    await showSetupOnFirstUse(globalState, workspaceState)
+    expect(mockedAskShowOrCloseOrNever).not.toHaveBeenCalled()
+  })
+
+  it('should not ask to show the setup page when the workspace state has data even if the global install key is not set', async () => {
+    const workspaceState = buildMockMemento({
+      [PersistenceKey.PLOT_HEIGHT + 'root1']: DEFAULT_HEIGHT
+    })
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
+
     expect(mockedAskShowOrCloseOrNever).not.toHaveBeenCalled()
   })
 
   it('should not ask to show the setup page when the user has set a config option', async () => {
     mockedGetConfigValue.mockReturnValueOnce(true)
-    await showSetupOnFirstUse(true)
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
     expect(mockedAskShowOrCloseOrNever).not.toHaveBeenCalled()
     expect(mockedGetConfigValue).toHaveBeenCalledTimes(1)
     expect(mockedGetConfigValue).toHaveBeenCalledWith(
@@ -47,7 +82,10 @@ describe('showSetupOnFirstUse', () => {
 
   it('should set the config option if the user responds with never', async () => {
     mockedAskShowOrCloseOrNever.mockResolvedValueOnce(Response.NEVER)
-    await showSetupOnFirstUse(true)
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
 
     expect(mockedSetConfigValue).toHaveBeenCalledTimes(1)
     expect(mockedSetConfigValue).toHaveBeenCalledWith(
@@ -58,7 +96,10 @@ describe('showSetupOnFirstUse', () => {
 
   it('should show the setup page if the user responds with show', async () => {
     mockedAskShowOrCloseOrNever.mockResolvedValueOnce(Response.SHOW)
-    await showSetupOnFirstUse(true)
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
 
     expect(mockedSetConfigValue).not.toHaveBeenCalled()
     expect(mockedExecuteCommand).toHaveBeenCalledWith(
@@ -68,15 +109,21 @@ describe('showSetupOnFirstUse', () => {
 
   it('should take no action if the user closes the dialog', async () => {
     mockedAskShowOrCloseOrNever.mockResolvedValueOnce(undefined)
-    await showSetupOnFirstUse(true)
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
 
     expect(mockedSetConfigValue).not.toHaveBeenCalled()
     expect(mockedExecuteCommand).not.toHaveBeenCalled()
   })
 
   it('should take no action if the user respond with close', async () => {
     mockedAskShowOrCloseOrNever.mockResolvedValueOnce(Response.CLOSE)
-    await showSetupOnFirstUse(true)
+    const workspaceState = buildMockMemento()
+    const globalState = buildMockMemento()
+
+    await showSetupOnFirstUse(globalState, workspaceState)
 
     expect(mockedSetConfigValue).not.toHaveBeenCalled()
     expect(mockedExecuteCommand).not.toHaveBeenCalled()

extension/src/setup/util.ts

@@ -1,12 +1,12 @@
-import { commands } from 'vscode'
+/* eslint-disable import/no-unused-modules */
+import { Memento, commands } from 'vscode'
 import { RegisteredCommands } from '../commands/external'
 import { ConfigKey, getConfigValue, setUserConfigValue } from '../vscode/config'
 import { Response } from '../vscode/response'
 import { Toast } from '../vscode/toast'
+import { GlobalPersistenceKey } from '../persistence/constants'
 
-export const showSetupOnFirstUse = async (
-  isNewAppInstall: boolean
-): Promise<void> => {
+const showSetupToast = async (isNewAppInstall: boolean): Promise<void> => {
   if (
     !isNewAppInstall ||
     getConfigValue<boolean>(ConfigKey.DO_NOT_SHOW_SETUP_AFTER_INSTALL)
@@ -25,3 +25,20 @@ export const showSetupOnFirstUse = async (
     void setUserConfigValue(ConfigKey.DO_NOT_SHOW_SETUP_AFTER_INSTALL, true)
   }
 }
+
+export const showSetupOnFirstUse = (
+  globalState: Memento,
+  workspaceState: Memento
+): Promise<void> | undefined => {
+  const installed = globalState.get(GlobalPersistenceKey.INSTALLED, false)
+
+  if (installed) {
+    return
+  }
+
+  void globalState.update(GlobalPersistenceKey.INSTALLED, true)
+
+  // old users won't have the installedKey even if it's not a new install
+  const workspaceKeys = workspaceState.keys()
+  return showSetupToast(workspaceKeys.length === 0)
+}