diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts index 977fbab88d..610ace8f66 100644 --- a/extension/src/fileSystem/index.test.ts +++ b/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 diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts index 5d5a607bd6..dc4c5fe39a 100644 --- a/extension/src/fileSystem/index.ts +++ b/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 diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index ece82cc470..1cbcacb3c9 100644 --- a/extension/src/pipeline/index.ts +++ b/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 }) => { diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts index 82613e1ef3..4f485c25ce 100644 --- a/extension/src/test/suite/pipeline/index.test.ts +++ b/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 () => {