Add "Add Top-level Plot" command to plots webview (#4664)

extension/src/pipeline/register.ts

@@ -1,6 +1,7 @@
 import { WorkspacePipeline } from './workspace'
 import { RegisteredCommands } from '../commands/external'
 import { InternalCommands } from '../commands/internal'
+import { Context, getDvcRootFromContext } from '../vscode/context'
 
 export const registerPipelineCommands = (
   pipelines: WorkspacePipeline,
@@ -13,6 +14,7 @@ export const registerPipelineCommands = (
 
   internalCommands.registerExternalCommand(
     RegisteredCommands.PIPELINE_ADD_PLOT,
-    () => pipelines.addTopLevelPlot()
+    (context: Context) =>
+      pipelines.addTopLevelPlot(getDvcRootFromContext(context))
   )
 }

extension/src/pipeline/workspace.ts

@@ -63,8 +63,8 @@ export class WorkspacePipeline extends BaseWorkspace<Pipeline> {
     )
   }
 
-  public async addTopLevelPlot() {
-    const cwd = await this.getCwd()
+  public async addTopLevelPlot(overrideRoot?: string) {
+    const cwd = overrideRoot || (await this.getCwd())
 
     if (!cwd) {
       return

extension/src/plots/webview/messages.ts

@@ -76,6 +76,11 @@ export class WebviewMessages {
           RegisteredCommands.PLOTS_CUSTOM_ADD,
           this.dvcRoot
         )
+      case MessageFromWebviewType.ADD_PIPELINE_PLOT:
+        return commands.executeCommand(
+          RegisteredCommands.PIPELINE_ADD_PLOT,
+          this.dvcRoot
+        )
       case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_CSV:
         return this.exportPlotDataAsCsv(message.payload)
       case MessageFromWebviewType.EXPORT_PLOT_DATA_AS_TSV:

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

@@ -1209,6 +1209,24 @@ suite('Plots Test Suite', () => {
       )
     })
 
+    it('should handle an add pipeline plot message from the webview', async () => {
+      const { mockMessageReceived } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+
+      const executeCommandSpy = spy(commands, 'executeCommand')
+
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.ADD_PIPELINE_PLOT
+      })
+
+      expect(executeCommandSpy).to.be.calledWithExactly(
+        RegisteredCommands.PIPELINE_ADD_PLOT,
+        dvcDemoPath
+      )
+    })
+
     it('should handle the CLI throwing an error', async () => {
       const { data, errorsModel, mockPlotsDiff, plots, plotsModel } =
         await buildPlots({ disposer: disposable, plotsDiff: plotsDiffFixture })

extension/src/webview/contract.ts

@@ -17,6 +17,7 @@ export enum MessageFromWebviewType {
   APPLY_EXPERIMENT_TO_WORKSPACE = 'apply-experiment-to-workspace',
   ADD_STARRED_EXPERIMENT_FILTER = 'add-starred-experiment-filter',
   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',
@@ -110,6 +111,9 @@ export type MessageFromWebview =
   | {
       type: MessageFromWebviewType.ADD_CUSTOM_PLOT
     }
+  | {
+      type: MessageFromWebviewType.ADD_PIPELINE_PLOT
+    }
   | {
       type: MessageFromWebviewType.COPY_TO_CLIPBOARD
       payload: string

webview/src/plots/components/App.test.tsx

@@ -336,18 +336,18 @@ describe('App', () => {
     expect(loading).toHaveLength(3)
   })
 
-  it('should render only get started (buttons: add plots, add experiments, add custom plots) when there are some selected exps, all unselected plots, and no custom plots', async () => {
+  it('should render only get started (buttons: select plots, add experiments, add custom plots) when there are some selected exps, all unselected plots, and no custom plots', async () => {
     renderAppWithOptionalData({
       hasPlots: true,
       hasUnselectedPlots: true,
       selectedRevisions: [{} as Revision]
     })
     const addExperimentsButton = await screen.findByText('Add Experiments')
-    const addPlotsButton = await screen.findByText('Add Plots')
+    const selectPlotsButton = await screen.findByText('Select Plots')
     const addCustomPlotsButton = await screen.findByText('Add Custom Plot')
 
     expect(addExperimentsButton).toBeInTheDocument()
-    expect(addPlotsButton).toBeInTheDocument()
+    expect(selectPlotsButton).toBeInTheDocument()
     expect(addCustomPlotsButton).toBeInTheDocument()
     expect(screen.queryByTestId('section-container')).not.toBeInTheDocument()
 
@@ -361,7 +361,7 @@ describe('App', () => {
 
     mockPostMessage.mockReset()
 
-    fireEvent.click(addPlotsButton)
+    fireEvent.click(selectPlotsButton)
 
     expect(mockPostMessage).toHaveBeenCalledWith({
       type: MessageFromWebviewType.SELECT_PLOTS
@@ -427,20 +427,20 @@ describe('App', () => {
     mockPostMessage.mockReset()
   })
 
-  it('should render get started (buttons: add plots, add experiments) and custom section when there are some selected exps, all unselected plots, and added custom plots', async () => {
+  it('should render get started (buttons: select plots, add experiments) and custom section when there are some selected exps, all unselected plots, and added custom plots', async () => {
     renderAppWithOptionalData({
       custom: customPlotsFixture,
       hasPlots: true,
       hasUnselectedPlots: true,
       selectedRevisions: [{} as Revision]
     })
     const addExperimentsButton = await screen.findByText('Add Experiments')
-    const addPlotsButton = await screen.findByText('Add Plots')
+    const selectPlotsButton = await screen.findByText('Select Plots')
     const addCustomPlotsButton = screen.queryByText('Add Custom Plot')
     const customSection = await screen.findByTestId('section-container')
 
     expect(addExperimentsButton).toBeInTheDocument()
-    expect(addPlotsButton).toBeInTheDocument()
+    expect(selectPlotsButton).toBeInTheDocument()
     expect(addCustomPlotsButton).not.toBeInTheDocument()
     expect(customSection).toBeInTheDocument()
 
@@ -454,29 +454,29 @@ describe('App', () => {
 
     mockPostMessage.mockReset()
 
-    fireEvent.click(addPlotsButton)
+    fireEvent.click(selectPlotsButton)
 
     expect(mockPostMessage).toHaveBeenCalledWith({
       type: MessageFromWebviewType.SELECT_PLOTS
     })
     mockPostMessage.mockReset()
   })
 
-  it('should render only get started (buttons: add experiments, add custom plots) when there are no selected exps and no custom plots', async () => {
+  it('should render only get started (buttons: add experiments, add custom plots, add plots) when there are no selected exps and no custom plots', async () => {
     renderAppWithOptionalData({
       custom: null,
       hasPlots: true,
       hasUnselectedPlots: false,
       selectedRevisions: undefined
     })
     const addExperimentsButton = await screen.findByText('Add Experiments')
-    const addPlotsButton = screen.queryByText('Add Plots')
+    const addPlotsButton = await screen.findByText('Add Plot')
     const addCustomPlotsButton = await screen.findByText('Add Custom Plot')
     const customSection = screen.queryByTestId('section-container')
 
     expect(addExperimentsButton).toBeInTheDocument()
     expect(addCustomPlotsButton).toBeInTheDocument()
-    expect(addPlotsButton).not.toBeInTheDocument()
+    expect(addPlotsButton).toBeInTheDocument()
     expect(customSection).not.toBeInTheDocument()
 
     mockPostMessage.mockReset()
@@ -492,23 +492,29 @@ describe('App', () => {
       type: MessageFromWebviewType.ADD_CUSTOM_PLOT
     })
     mockPostMessage.mockReset()
+
+    fireEvent.click(addPlotsButton)
+
+    expect(mockPostMessage).toHaveBeenCalledWith({
+      type: MessageFromWebviewType.ADD_PIPELINE_PLOT
+    })
   })
 
-  it('should render get started (buttons: add experiments) and custom section when there are no selected exps and added custom plots', async () => {
+  it('should render get started (buttons: add experiments, add plots) and custom section when there are no selected exps and added custom plots', async () => {
     renderAppWithOptionalData({
       custom: customPlotsFixture,
       hasPlots: true,
       hasUnselectedPlots: false,
       selectedRevisions: undefined
     })
     const addExperimentsButton = await screen.findByText('Add Experiments')
-    const addPlotsButton = screen.queryByText('Add Plots')
+    const addPlotsButton = await screen.findByText('Add Plot')
     const addCustomPlotsButton = screen.queryByText('Add Custom Plot')
     const customSection = await screen.findByTestId('section-container')
 
     expect(addExperimentsButton).toBeInTheDocument()
     expect(addCustomPlotsButton).not.toBeInTheDocument()
-    expect(addPlotsButton).not.toBeInTheDocument()
+    expect(addPlotsButton).toBeInTheDocument()
     expect(customSection).toBeInTheDocument()
 
     mockPostMessage.mockReset()
@@ -517,6 +523,13 @@ describe('App', () => {
     expect(mockPostMessage).toHaveBeenCalledWith({
       type: MessageFromWebviewType.SELECT_EXPERIMENTS
     })
+
+    mockPostMessage.mockReset()
+
+    fireEvent.click(addPlotsButton)
+    expect(mockPostMessage).toHaveBeenCalledWith({
+      type: MessageFromWebviewType.ADD_PIPELINE_PLOT
+    })
   })
 
   it('should render custom with "No Plots to Display" message when there is no custom plots data', () => {
@@ -763,16 +776,14 @@ describe('App', () => {
 
     const customSection = screen.getAllByTestId('section-container')[2]
 
-    expect(
-      within(customSection).getByLabelText('Add Plots')
-    ).toBeInTheDocument()
+    expect(within(customSection).getByLabelText('Add Plot')).toBeInTheDocument()
 
     sendSetDataMessage({
       custom: { ...customPlotsFixture, enablePlotCreation: false }
     })
 
     expect(
-      within(customSection).queryByLabelText('Add Plots')
+      within(customSection).queryByLabelText('Add Plot')
     ).not.toBeInTheDocument()
   })
 

webview/src/plots/components/PlotsContainer.tsx

@@ -74,7 +74,7 @@ export const PlotsContainer: React.FC<PlotsContainerProps> = ({
     menuItems.unshift({
       icon: Add,
       onClick: addPlotsButton.onClick,
-      tooltip: 'Add Plots'
+      tooltip: 'Add Plot'
     })
   }
 

webview/src/plots/components/emptyState/AddPlots.tsx

@@ -1,6 +1,7 @@
 import React from 'react'
 import {
   addCustomPlot,
+  addPipelinePlot,
   selectPlots,
   selectRevisions
 } from '../../util/messages'
@@ -19,12 +20,19 @@ export const AddPlots: React.FC<AddPlotsProps> = ({
     <p>No Plots to Display</p>
     <div>
       <StartButton onClick={selectRevisions} text="Add Experiments" />
-      {hasUnselectedPlots && (
+      {hasUnselectedPlots ? (
         <StartButton
           isNested={true}
           appearance="secondary"
           onClick={selectPlots}
-          text="Add Plots"
+          text="Select Plots"
+        />
+      ) : (
+        <StartButton
+          isNested={true}
+          appearance="secondary"
+          onClick={addPipelinePlot}
+          text="Add Plot"
         />
       )}
       {!hasCustomPlots && (

webview/src/plots/components/emptyState/Welcome.tsx

@@ -1,11 +1,19 @@
 import React from 'react'
 import { StartButton } from '../../../shared/components/button/StartButton'
-import { selectRevisions } from '../../util/messages'
+import { addPipelinePlot, selectRevisions } from '../../util/messages'
 
 export const Welcome: React.FC = () => (
   <div>
     <p>No Plots Detected.</p>
-    <StartButton onClick={selectRevisions} text="Add Experiments" />
+    <div>
+      <StartButton onClick={addPipelinePlot} text="Add Plot" />
+      <StartButton
+        appearance="secondary"
+        onClick={selectRevisions}
+        text="Add Experiments"
+        isNested={true}
+      />
+    </div>
     <p>
       Learn how to{' '}
       <a href="https://dvc.org/doc/user-guide/experiment-management/visualizing-plots">

webview/src/plots/components/templatePlots/TemplatePlotsWrapper.tsx

@@ -5,6 +5,7 @@ import { TemplatePlots } from './TemplatePlots'
 import { changeSize } from './templatePlotsSlice'
 import { PlotsContainer } from '../PlotsContainer'
 import { PlotsState } from '../../store'
+import { addPipelinePlot } from '../../util/messages'
 
 export const TemplatePlotsWrapper: React.FC = () => {
   const { nbItemsPerRow, isCollapsed, height, hasItems } = useSelector(
@@ -20,6 +21,7 @@ export const TemplatePlotsWrapper: React.FC = () => {
       changeSize={changeSize}
       hasItems={hasItems}
       height={height}
+      addPlotsButton={{ onClick: addPipelinePlot }}
     >
       <TemplatePlots />
     </PlotsContainer>

webview/src/plots/util/messages.ts

@@ -8,6 +8,12 @@ export const addCustomPlot = () =>
     type: MessageFromWebviewType.ADD_CUSTOM_PLOT
   })
 
+export const addPipelinePlot = () => {
+  sendMessage({
+    type: MessageFromWebviewType.ADD_PIPELINE_PLOT
+  })
+}
+
 export const exportPlotDataAsCsv = (id: string) => {
   sendMessage({
     payload: id,