Remove check for stages in plot wizard (#4904)

Author: julieg18

extension/src/fileSystem/index.test.ts

@@ -640,6 +640,29 @@ describe('addPlotToDvcYamlFile', () => {
     '      y:',
     '        data.json: accuracy'
   ]
+
+  it('should add a plots list with the new plot if the dvc.yaml file does not exist', () => {
+    const mockPlotYamlContent = ['', 'plots:', ...mockNewPlotLines, ''].join(
+      '\n'
+    )
+    mockedReadFileSync.mockReturnValueOnce('')
+    mockedReadFileSync.mockReturnValueOnce('')
+
+    addPlotToDvcYamlFile('/', {
+      template: 'simple',
+      title: 'Simple Plot',
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
+    })
+
+    expect(mockedEnsureFileSync).toHaveBeenCalledWith('//dvc.yaml')
+    expect(mockedOpenTextDocument).toHaveBeenCalledTimes(1)
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockPlotYamlContent
+    )
+  })
+
   it('should add a plots list with the new plot if the dvc.yaml file has no plots', () => {
     const mockDvcYamlContent = mockStagesLines.join('\n')
     const mockPlotYamlContent = ['', 'plots:', ...mockNewPlotLines, ''].join(
@@ -655,6 +678,7 @@ describe('addPlotToDvcYamlFile', () => {
       y: { 'data.json': ['accuracy'] }
     })
 
+    expect(mockedOpenTextDocument).toHaveBeenCalledTimes(1)
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
       '//dvc.yaml',
       mockDvcYamlContent + mockPlotYamlContent
@@ -684,6 +708,7 @@ describe('addPlotToDvcYamlFile', () => {
       y: { 'acc.json': ['accuracy'] }
     })
 
+    expect(mockedOpenTextDocument).toHaveBeenCalledTimes(1)
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
       '//dvc.yaml',
       mockDvcYamlContent + mockPlotYamlContent
@@ -714,6 +739,7 @@ describe('addPlotToDvcYamlFile', () => {
       y: { 'data.json': ['accuracy', 'epochs'] }
     })
 
+    expect(mockedOpenTextDocument).toHaveBeenCalledTimes(1)
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
       '//dvc.yaml',
       mockDvcYamlContent + mockPlotYamlContent
@@ -735,6 +761,7 @@ describe('addPlotToDvcYamlFile', () => {
 
     mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
 
+    expect(mockedOpenTextDocument).toHaveBeenCalledTimes(1)
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
       '//dvc.yaml',
       mockDvcYamlContent.join('\n')
@@ -756,6 +783,7 @@ describe('addPlotToDvcYamlFile', () => {
       y: { 'data.json': ['accuracy'] }
     })
 
+    expect(mockedOpenTextDocument).toHaveBeenCalledTimes(1)
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
       '//dvc.yaml',
       mockDvcYamlContent + mockPlotYamlContent
@@ -794,6 +822,7 @@ describe('addPlotToDvcYamlFile', () => {
       y: { 'data.json': ['accuracy'] }
     })
 
+    expect(mockedOpenTextDocument).toHaveBeenCalledTimes(1)
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
       '//dvc.yaml',
       mockDvcYamlContent + mockPlotYamlContent

extension/src/fileSystem/index.ts

@@ -262,6 +262,7 @@ const getPlotsYaml = (plotObj: PlotConfigData, indentSearchLines: string[]) => {
 
 export const addPlotToDvcYamlFile = (cwd: string, plotObj: PlotConfigData) => {
   const dvcYamlFile = `${cwd}/dvc.yaml`
+  ensureFileSync(dvcYamlFile)
   const dvcYamlDoc = loadYamlAsDoc(dvcYamlFile)
 
   if (!dvcYamlDoc) {
@@ -276,8 +277,9 @@ export const addPlotToDvcYamlFile = (cwd: string, plotObj: PlotConfigData) => {
   if (!plots?.range) {
     const plotYaml = getPlotsYaml(plotObj, dvcYamlLines)
     dvcYamlLines.push(...plotYaml)
-    writeFileSync(dvcYamlFile, dvcYamlLines.join('\n'))
-    return
+
+    void openFileInEditor(dvcYamlFile)
+    return writeFileSync(dvcYamlFile, dvcYamlLines.join('\n'))
   }
 
   const plotsEndPos = lineCounter.linePos(plots.range[2]).line

extension/src/pipeline/index.ts

@@ -83,28 +83,8 @@ export class Pipeline extends DeferredDisposable {
   }
 
   public async getCwd() {
-    const focusedPipeline = this.getFocusedPipeline()
-    if (focusedPipeline) {
-      return focusedPipeline
-    }
-
     await this.checkOrAddPipeline()
-
-    const pipelines = this.model.getPipelines()
-    if (!pipelines?.size) {
-      return
-    }
-    if (pipelines.has(this.dvcRoot)) {
-      return this.dvcRoot
-    }
-    if (pipelines.size === 1) {
-      return [...pipelines][0]
-    }
-
-    return quickPickOne(
-      [...pipelines],
-      'Select a Pipeline to Run Command Against'
-    )
+    return this.findPipelineCwd()
   }
 
   public async checkOrAddPipeline() {
@@ -158,11 +138,7 @@ export class Pipeline extends DeferredDisposable {
   }
 
   public async addDataSeriesPlot() {
-    const cwd = await this.getCwd()
-
-    if (!cwd) {
-      return
-    }
+    const cwd = (await this.findPipelineCwd()) || this.dvcRoot
 
     const plotConfiguration = await pickPlotConfiguration(cwd)
 
@@ -177,6 +153,29 @@ export class Pipeline extends DeferredDisposable {
     return appendFileSync(join(this.dvcRoot, TEMP_DAG_FILE), '\n')
   }
 
+  private findPipelineCwd() {
+    const focusedPipeline = this.getFocusedPipeline()
+    if (focusedPipeline) {
+      return focusedPipeline
+    }
+
+    const pipelines = this.model.getPipelines()
+    if (!pipelines?.size) {
+      return
+    }
+    if (pipelines.has(this.dvcRoot)) {
+      return this.dvcRoot
+    }
+    if (pipelines.size === 1) {
+      return [...pipelines][0]
+    }
+
+    return quickPickOne(
+      [...pipelines],
+      'Select a Pipeline to Run Command Against'
+    )
+  }
+
   private async initialize() {
     this.dispose.track(
       this.data.onDidUpdate(({ dag, stages }) => {

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

@@ -323,10 +323,9 @@ suite('Pipeline Test Suite', () => {
     })
   })
 
-  it('should add a data series plot', async () => {
+  it('should add a data series plot when a dvc.yaml file exists', async () => {
     const { pipeline } = buildPipeline({
-      disposer: disposable,
-      dvcRoot: dvcDemoPath
+      disposer: disposable
     })
     const mockPickPlotConfiguration = stub(
       PipelineQuickPick,
@@ -340,7 +339,7 @@ suite('Pipeline Test Suite', () => {
 
     await pipeline.addDataSeriesPlot()
 
-    expect(mockPickPlotConfiguration).to.be.calledOnce
+    expect(mockPickPlotConfiguration).to.be.calledOnceWithExactly(dvcDemoPath)
     expect(mockAddPlotToDvcFile).not.to.be.called
 
     const mockPlotConfig: PipelineQuickPick.PlotConfigData = {
@@ -354,8 +353,43 @@ suite('Pipeline Test Suite', () => {
 
     await pipeline.addDataSeriesPlot()
 
-    expect(mockPickPlotConfiguration).to.be.calledTwice
-    expect(mockAddPlotToDvcFile).to.be.called
+    expect(mockPickPlotConfiguration).to.be.calledWithExactly(dvcDemoPath)
+    expect(mockAddPlotToDvcFile).to.be.calledOnceWithExactly(
+      dvcDemoPath,
+      mockPlotConfig
+    )
+  })
+
+  it('should add a data series plot without trying to add a missing dvc.yaml file or stage', async () => {
+    const { pipeline } = buildPipeline({
+      disposer: disposable,
+      dvcYamls: []
+    })
+    const mockPickPlotConfiguration = stub(
+      PipelineQuickPick,
+      'pickPlotConfiguration'
+    )
+    const mockAddPlotToDvcFile = stub(FileSystem, 'addPlotToDvcYamlFile')
+    const mockCheckOrAddPipeline = stub(pipeline, 'checkOrAddPipeline')
+    const mockPlotConfig: PipelineQuickPick.PlotConfigData = {
+      template: 'simple',
+      title: 'Great Plot Name',
+      x: { 'results.json': ['step'] },
+      y: { 'results.json': ['acc'] }
+    }
+
+    mockPickPlotConfiguration.onFirstCall().resolves(mockPlotConfig)
+
+    await pipeline.isReady()
+    await pipeline.addDataSeriesPlot()
+
+    expect(mockCheckOrAddPipeline, 'should not check for a pipeline stage').not
+      .to.be.called
+    expect(mockPickPlotConfiguration).to.be.calledOnceWithExactly(dvcDemoPath)
+    expect(mockAddPlotToDvcFile).to.be.calledOnceWithExactly(
+      dvcDemoPath,
+      mockPlotConfig
+    )
   })
 
   it('should set the appropriate context value when a dvc.yaml is open in the active editor', async () => {