Merge create plot commands

extension/package.json

@@ -425,7 +425,7 @@
       },
       {
         "title": "Add Plot",
-        "command": "dvc.addTopLevelPlot",
+        "command": "dvc.addPlot",
         "category": "DVC",
         "icon": "$(add)"
       },
@@ -601,11 +601,6 @@
         "category": "DVC",
         "icon": "$(refresh)"
       },
-      {
-        "title": "Add Custom Plot",
-        "command": "dvc.views.plots.addCustomPlot",
-        "category": "DVC"
-      },
       {
         "title": "Remove Custom Plot(s)",
         "command": "dvc.views.plots.removeCustomPlots",
@@ -910,7 +905,7 @@
           "when": "dvc.commands.available && dvc.project.available"
         },
         {
-          "command": "dvc.addTopLevelPlot",
+          "command": "dvc.addPlot",
           "when": "dvc.commands.available && dvc.project.available"
         },
         {
@@ -1029,10 +1024,6 @@
           "command": "dvc.views.plots.refreshPlots",
           "when": "dvc.commands.available && dvc.project.available"
         },
-        {
-          "command": "dvc.views.plots.addCustomPlot",
-          "when": "dvc.commands.available && dvc.project.available"
-        },
         {
           "command": "dvc.views.plots.removeCustomPlots",
           "when": "dvc.commands.available && dvc.project.available"
@@ -1425,7 +1416,7 @@
           "group": "navigation@3"
         },
         {
-          "command": "dvc.addTopLevelPlot",
+          "command": "dvc.addPlot",
           "when": "view == dvc.views.plotsPathsTree",
           "group": "navigation@0"
         },

extension/src/commands/external.ts

@@ -81,6 +81,8 @@ export enum RegisteredCommands {
   PLOTS_CUSTOM_ADD = 'dvc.views.plots.addCustomPlot',
   PLOTS_CUSTOM_REMOVE = 'dvc.views.plots.removeCustomPlots',
 
+  ADD_PLOT = 'dvc.addPlot',
+
   EXPERIMENT_AND_PLOTS_SHOW = 'dvc.showExperimentsAndPlots',
 
   EXTENSION_CHECK_CLI_COMPATIBLE = 'dvc.checkCLICompatible',

extension/src/commands/util.test.ts

@@ -0,0 +1,66 @@
+import { commands } from 'vscode'
+import { RegisteredCommands } from './external'
+import { addPlotCommand } from './util'
+import { quickPickValue } from '../vscode/quickPick'
+import { Title } from '../vscode/title'
+
+jest.mock('../vscode/quickPick')
+
+const mockedQuickPickValue = jest.mocked(quickPickValue)
+const mockedCommands = jest.mocked(commands)
+const mockedExecuteCommand = jest.fn()
+mockedCommands.executeCommand = mockedExecuteCommand
+
+beforeEach(() => {
+  jest.resetAllMocks()
+})
+
+describe('addPlotCommand', () => {
+  it('should call no add commands if no plots types are selected', async () => {
+    mockedQuickPickValue.mockResolvedValueOnce(undefined)
+
+    await addPlotCommand(undefined)
+
+    expect(mockedQuickPickValue).toHaveBeenCalledWith(
+      [
+        {
+          description: 'Create a dvc.yaml plot based off a chosen data file',
+          label: 'Top-Level',
+          value: 'top-level'
+        },
+        {
+          description:
+            'Create an extension-only plot based off a chosen metric and param',
+          label: 'Custom',
+          value: 'custom'
+        }
+      ],
+      {
+        title: Title.SELECT_PLOT_TYPE
+      }
+    )
+    expect(mockedExecuteCommand).not.toHaveBeenCalled()
+  })
+
+  it('should call the add top-level plot command if the top-level type is selected', async () => {
+    mockedQuickPickValue.mockResolvedValueOnce('top-level')
+
+    await addPlotCommand(undefined)
+
+    expect(mockedExecuteCommand).toHaveBeenCalledWith(
+      RegisteredCommands.PIPELINE_ADD_PLOT,
+      undefined
+    )
+  })
+
+  it('should call the add custom plot command if the custom type is selected', async () => {
+    mockedQuickPickValue.mockResolvedValueOnce('custom')
+
+    await addPlotCommand('cwd')
+
+    expect(mockedExecuteCommand).toHaveBeenCalledWith(
+      RegisteredCommands.PLOTS_CUSTOM_ADD,
+      'cwd'
+    )
+  })
+})

extension/src/commands/util.ts

@@ -2,6 +2,8 @@ import { commands } from 'vscode'
 import { RegisteredCommands } from './external'
 import { Setup } from '../setup'
 import { Context } from '../vscode/context'
+import { quickPickValue } from '../vscode/quickPick'
+import { Title } from '../vscode/title'
 
 export const showSetupOrExecuteCommand =
   <T>(setup: Setup, callback: (context: Context) => Promise<T | undefined>) =>
@@ -18,3 +20,40 @@ export const showSetupOrExecuteCommand =
 
     return callback(context)
   }
+
+const enum PLOT_TYPE {
+  CUSTOM = 'custom',
+  TOP_LEVEL = 'top-level'
+}
+
+export const addPlotCommand = async (context: Context) => {
+  const result: PLOT_TYPE | undefined = await quickPickValue(
+    [
+      {
+        description: 'Create a dvc.yaml plot based off a chosen data file',
+        label: 'Top-Level',
+        value: PLOT_TYPE.TOP_LEVEL
+      },
+      {
+        description:
+          'Create an extension-only plot based off a chosen metric and param',
+        label: 'Custom',
+        value: PLOT_TYPE.CUSTOM
+      }
+    ],
+    {
+      title: Title.SELECT_PLOT_TYPE
+    }
+  )
+
+  if (result === PLOT_TYPE.CUSTOM) {
+    return commands.executeCommand(RegisteredCommands.PLOTS_CUSTOM_ADD, context)
+  }
+
+  if (result === PLOT_TYPE.TOP_LEVEL) {
+    return commands.executeCommand(
+      RegisteredCommands.PIPELINE_ADD_PLOT,
+      context
+    )
+  }
+}

extension/src/extension.ts

@@ -49,7 +49,7 @@ import { DvcViewer } from './cli/dvc/viewer'
 import { registerSetupCommands } from './setup/commands/register'
 import { Status } from './status'
 import { registerPersistenceCommands } from './persistence/register'
-import { showSetupOrExecuteCommand } from './commands/util'
+import { addPlotCommand, showSetupOrExecuteCommand } from './commands/util'
 import { WorkspacePipeline } from './pipeline/workspace'
 import { registerPipelineCommands } from './pipeline/register'
 
@@ -193,6 +193,10 @@ class Extension extends Disposable {
     )
     registerPipelineCommands(this.pipelines, this.internalCommands)
     registerPlotsCommands(this.plots, this.internalCommands, this.setup)
+    this.internalCommands.registerExternalCommand(
+      RegisteredCommands.ADD_PLOT,
+      addPlotCommand
+    )
     registerSetupCommands(this.setup, this.internalCommands)
     this.internalCommands.registerExternalCommand(
       RegisteredCommands.EXPERIMENT_AND_PLOTS_SHOW,

extension/src/plots/webview/messages.ts

@@ -71,14 +71,14 @@ export class WebviewMessages {
 
   public handleMessageFromWebview(message: MessageFromWebview) {
     switch (message.type) {
-      case MessageFromWebviewType.ADD_CUSTOM_PLOT:
+      case MessageFromWebviewType.ADD_PLOT:
         return commands.executeCommand(
-          RegisteredCommands.PLOTS_CUSTOM_ADD,
+          RegisteredCommands.ADD_PLOT,
           this.dvcRoot
         )
-      case MessageFromWebviewType.ADD_PIPELINE_PLOT:
+      case MessageFromWebviewType.ADD_CUSTOM_PLOT:
         return commands.executeCommand(
-          RegisteredCommands.PIPELINE_ADD_PLOT,
+          RegisteredCommands.PLOTS_CUSTOM_ADD,
           this.dvcRoot
         )
       case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV:

extension/src/telemetry/constants.ts

@@ -204,6 +204,8 @@ export interface IEventNamePropertyMapping {
   [EventName.PLOTS_CUSTOM_ADD]: undefined
   [EventName.PLOTS_CUSTOM_REMOVE]: undefined
 
+  [EventName.ADD_PLOT]: undefined
+
   [EventName.ADD_TARGET]: undefined
   [EventName.CHECKOUT_TARGET]: undefined
   [EventName.CHECKOUT]: undefined

extension/src/test/e2e/util.ts

@@ -185,7 +185,9 @@ export const closeAllEditors = async (): Promise<void> => {
 
 export const createCustomPlot = async (): Promise<void> => {
   const workbench = await browser.getWorkbench()
-  const addCustomPlot = await workbench.executeCommand('DVC: Add Custom Plot')
+  const addCustomPlot = await workbench.executeCommand('DVC: Add Plot')
+  await browser.waitUntil(() => addCustomPlot.elem.isDisplayed())
+  await browser.keys(['ArrowDown', 'Enter'])
   await browser.waitUntil(() => addCustomPlot.elem.isDisplayed())
   await browser.keys('Enter')
   await browser.waitUntil(() => addCustomPlot.elem.isDisplayed())

extension/src/test/suite/plots/index.test.ts

@@ -1209,7 +1209,7 @@ suite('Plots Test Suite', () => {
       )
     })
 
-    it('should handle an add pipeline plot message from the webview', async () => {
+    it('should handle an add plot message from the webview', async () => {
       const { mockMessageReceived } = await buildPlotsWebview({
         disposer: disposable,
         plotsDiff: plotsDiffFixture
@@ -1218,11 +1218,11 @@ suite('Plots Test Suite', () => {
       const executeCommandSpy = spy(commands, 'executeCommand')
 
       mockMessageReceived.fire({
-        type: MessageFromWebviewType.ADD_PIPELINE_PLOT
+        type: MessageFromWebviewType.ADD_PLOT
       })
 
       expect(executeCommandSpy).to.be.calledWithExactly(
-        RegisteredCommands.PIPELINE_ADD_PLOT,
+        RegisteredCommands.ADD_PLOT,
         dvcDemoPath
       )
     })

extension/src/vscode/title.ts

@@ -41,6 +41,7 @@ export enum Title {
   SELECT_SORTS_TO_REMOVE = 'Select Sort(s) to Remove',
   SELECT_TRAINING_SCRIPT = 'Select your Training Script',
   SELECT_PLOT_DATA = 'Select Data to Plot',
+  SELECT_PLOT_TYPE = 'Select Plot Type',
   SETUP_WORKSPACE = 'Setup the Workspace',
   SET_EXPERIMENTS_HEADER_HEIGHT = 'Set the Maximum Experiment Table Header Height',
   SET_REMOTE_AS_DEFAULT = 'Set Default Remote',

extension/src/webview/contract.ts

@@ -16,8 +16,8 @@ export enum MessageFromWebviewType {
   ADD_CONFIGURATION = 'add-configuration',
   APPLY_EXPERIMENT_TO_WORKSPACE = 'apply-experiment-to-workspace',
   ADD_STARRED_EXPERIMENT_FILTER = 'add-starred-experiment-filter',
+  ADD_PLOT = 'add-plot',
   ADD_CUSTOM_PLOT = 'add-custom-plot',
-  ADD_PIPELINE_PLOT = 'add-pipeline-plot',
   CREATE_BRANCH_FROM_EXPERIMENT = 'create-branch-from-experiment',
   COPY_TO_CLIPBOARD = 'copy-to-clipboard',
   COPY_STUDIO_LINK = 'copy-studio-link',
@@ -112,7 +112,7 @@ export type MessageFromWebview =
       type: MessageFromWebviewType.ADD_CUSTOM_PLOT
     }
   | {
-      type: MessageFromWebviewType.ADD_PIPELINE_PLOT
+      type: MessageFromWebviewType.ADD_PLOT
     }
   | {
       type: MessageFromWebviewType.COPY_TO_CLIPBOARD