From bf393d367c2e81deb35145d171d29dbe9cc8d212 Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Wed, 12 Jul 2023 10:57:53 +1000 Subject: [PATCH 01/12] add pipeline model for stages data --- extension/src/cli/dvc/reader.ts | 4 +- extension/src/data/index.ts | 2 +- extension/src/pipeline/data.ts | 15 ++++--- extension/src/pipeline/index.ts | 10 +++-- extension/src/pipeline/model/collect.test.ts | 44 +++++++++++++++++++ extension/src/pipeline/model/collect.ts | 31 +++++++++++++ extension/src/pipeline/model/index.ts | 26 +++++++++++ .../src/test/suite/experiments/index.test.ts | 22 +++++----- .../experiments/model/filterBy/tree.test.ts | 4 +- .../test/suite/experiments/model/tree.test.ts | 6 +-- .../test/suite/experiments/workspace.test.ts | 28 ++++++------ extension/src/test/suite/extension.test.ts | 2 +- 12 files changed, 151 insertions(+), 43 deletions(-) create mode 100644 extension/src/pipeline/model/collect.test.ts create mode 100644 extension/src/pipeline/model/collect.ts create mode 100644 extension/src/pipeline/model/index.ts diff --git a/extension/src/cli/dvc/reader.ts b/extension/src/cli/dvc/reader.ts index 54523f80f2..54c6965fb0 100644 --- a/extension/src/cli/dvc/reader.ts +++ b/extension/src/cli/dvc/reader.ts @@ -38,7 +38,7 @@ export const autoRegisteredCommands = { GLOBAL_VERSION: 'globalVersion', PLOTS_DIFF: 'plotsDiff', ROOT: 'root', - STAGE_LIST: 'listStages', + STAGE_LIST: 'stageList', VERSION: 'version' } as const @@ -144,7 +144,7 @@ export class DvcReader extends DvcCli { return this.executeProcess(options) } - public async listStages(cwd: string): Promise { + public async stageList(cwd: string): Promise { try { return await this.executeDvcProcess(cwd, Command.STAGE, SubCommand.LIST) } catch {} diff --git a/extension/src/data/index.ts b/extension/src/data/index.ts index 0ba68dfac9..6e3b762d33 100644 --- a/extension/src/data/index.ts +++ b/extension/src/data/index.ts @@ -20,7 +20,7 @@ export abstract class BaseData< T extends | { data: PlotsOutputOrError; revs: string[] } | ExperimentsOutput - | string + | { dag: string; stageList: string } > extends DeferredDisposable { public readonly onDidUpdate: Event public readonly onDidChangeDvcYaml: Event diff --git a/extension/src/pipeline/data.ts b/extension/src/pipeline/data.ts index 23265f0fd3..bf6dd09c6c 100644 --- a/extension/src/pipeline/data.ts +++ b/extension/src/pipeline/data.ts @@ -1,7 +1,7 @@ import { AvailableCommands, InternalCommands } from '../commands/internal' import { BaseData } from '../data' -export class PipelineData extends BaseData { +export class PipelineData extends BaseData<{ dag: string; stageList: string }> { constructor(dvcRoot: string, internalCommands: InternalCommands) { super( dvcRoot, @@ -18,10 +18,13 @@ export class PipelineData extends BaseData { } public async update(): Promise { - const dag = await this.internalCommands.executeCommand( - AvailableCommands.DAG, - this.dvcRoot - ) - return this.notifyChanged(dag) + const [dag, stageList] = await Promise.all([ + this.internalCommands.executeCommand(AvailableCommands.DAG, this.dvcRoot), + this.internalCommands.executeCommand( + AvailableCommands.STAGE_LIST, + this.dvcRoot + ) + ]) + return this.notifyChanged({ dag, stageList }) } } diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index 134e0299a1..6d5bad767b 100644 --- a/extension/src/pipeline/index.ts +++ b/extension/src/pipeline/index.ts @@ -1,6 +1,7 @@ import { join } from 'path' import { appendFileSync, writeFileSync } from 'fs-extra' import { PipelineData } from './data' +import { PipelineModel } from './model' import { DeferredDisposable } from '../class/deferred' import { InternalCommands } from '../commands/internal' import { TEMP_DAG_FILE } from '../cli/dvc/constants' @@ -8,11 +9,13 @@ import { TEMP_DAG_FILE } from '../cli/dvc/constants' export class Pipeline extends DeferredDisposable { private readonly dvcRoot: string private readonly data: PipelineData + private readonly model: PipelineModel constructor(dvcRoot: string, internalCommands: InternalCommands) { super() this.dvcRoot = dvcRoot this.data = this.dispose.track(new PipelineData(dvcRoot, internalCommands)) + this.model = this.dispose.track(new PipelineModel(dvcRoot)) void this.initialize() } @@ -23,9 +26,10 @@ export class Pipeline extends DeferredDisposable { private async initialize() { this.dispose.track( - this.data.onDidUpdate(data => - writeFileSync(join(this.dvcRoot, TEMP_DAG_FILE), data) - ) + this.data.onDidUpdate(({ dag, stageList }) => { + writeFileSync(join(this.dvcRoot, TEMP_DAG_FILE), dag) + this.model.transformAndSet(stageList) + }) ) await this.data.isReady() diff --git a/extension/src/pipeline/model/collect.test.ts b/extension/src/pipeline/model/collect.test.ts new file mode 100644 index 0000000000..290c893ab7 --- /dev/null +++ b/extension/src/pipeline/model/collect.test.ts @@ -0,0 +1,44 @@ +import { join } from 'path' +import { collectPipelines, collectStages } from './collect' +import { dvcDemoPath } from '../../test/util' + +describe('collectStages', () => { + it('should handle a simple list of stages', () => { + const stages = collectStages( + 'data.dvc Outputs data\ntrain Outputs model.pt, training/plots, hist.csv; Reports training/metrics.json' + ) + expect(stages).toStrictEqual(['train']) + }) + + it('should handle a monorepos list of stages', () => { + const stages = collectStages( + `nested1/data/data.xml.dvc Outputs data.xml + nested1/dvc.yaml:prepare Outputs data/prepared + nested1/dvc.yaml:featurize Outputs data/features + nested1/dvc.yaml:train Outputs model.pkl + nested1/dvc.yaml:evaluate Outputs eval/importance.png, eval/live/plots, eval/prc; Reports eval/live/metri… + nested2/data/data.xml.dvc Outputs data.xml` + ) + expect(stages).toStrictEqual([ + 'nested1/dvc.yaml:prepare', + 'nested1/dvc.yaml:featurize', + 'nested1/dvc.yaml:train', + 'nested1/dvc.yaml:evaluate' + ]) + }) +}) + +describe('collectPipelines', () => { + it('should handle a list of stages', () => { + const pipelines = collectPipelines(dvcDemoPath, [ + 'train', + 'nested1/dvc.yaml:prepare', + 'nested1/dvc.yaml:featurize', + 'nested1/dvc.yaml:train', + 'nested1/dvc.yaml:evaluate' + ]) + expect(pipelines).toStrictEqual( + new Set([dvcDemoPath, join(dvcDemoPath, 'nested1')]) + ) + }) +}) diff --git a/extension/src/pipeline/model/collect.ts b/extension/src/pipeline/model/collect.ts new file mode 100644 index 0000000000..b0253d183a --- /dev/null +++ b/extension/src/pipeline/model/collect.ts @@ -0,0 +1,31 @@ +import { join } from 'path' +import { trimAndSplit } from '../../util/stdout' + +export const collectStages = (stageList: string): string[] => { + const stages: string[] = [] + for (const stageStr of trimAndSplit(stageList)) { + const stage = stageStr.split(/\s+/)?.[0]?.trim() + if (!stage || stage.endsWith('.dvc')) { + continue + } + stages.push(stage) + } + return stages +} + +export const collectPipelines = ( + dvcRoot: string, + stages: string[] +): Set => { + const dvcYamlIdentifier = '/dvc.yaml:' + const pipelines = new Set() + for (const stage of stages) { + if (!stage.includes(dvcYamlIdentifier)) { + pipelines.add(dvcRoot) + continue + } + const [pipeline] = stage.split(dvcYamlIdentifier) + pipelines.add(join(dvcRoot, ...pipeline.split('/'))) + } + return pipelines +} diff --git a/extension/src/pipeline/model/index.ts b/extension/src/pipeline/model/index.ts new file mode 100644 index 0000000000..205db14ddb --- /dev/null +++ b/extension/src/pipeline/model/index.ts @@ -0,0 +1,26 @@ +import { collectPipelines, collectStages } from './collect' +import { Disposable } from '../../class/dispose' + +export class PipelineModel extends Disposable { + private readonly dvcRoot: string + private stages: string[] = [] + private pipelines: Set = new Set() + + constructor(dvcRoot: string) { + super() + this.dvcRoot = dvcRoot + } + + public transformAndSet(stageList: string) { + this.stages = collectStages(stageList) + this.pipelines = collectPipelines(this.dvcRoot, this.stages) + } + + public hasStages() { + return this.stages.length > 0 + } + + public getPipelines() { + return this.pipelines + } +} diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index ab5bc6e7d2..39e4e93970 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -122,7 +122,7 @@ suite('Experiments Test Suite', () => { describe('showWebview', () => { it('should be able to make the experiment webview visible', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, messageSpy } = buildExperiments({ disposer: disposable @@ -176,7 +176,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasValidDvcYaml to false if there is an error getting stages and there is a dvc.yaml file', async () => { - stub(DvcReader.prototype, 'listStages').resolves(undefined) + stub(DvcReader.prototype, 'stageList').resolves(undefined) stub(FileSystem, 'hasDvcYamlFile').returns(true) const { experiments, messageSpy } = buildExperiments({ @@ -191,7 +191,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasValidDvcYaml to true if there is an error getting stages and there is no dvc.yaml file', async () => { - stub(DvcReader.prototype, 'listStages').resolves(undefined) + stub(DvcReader.prototype, 'stageList').resolves(undefined) stub(FileSystem, 'hasDvcYamlFile').returns(false) const { experiments, messageSpy } = buildExperiments({ @@ -208,7 +208,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasValidDvcYaml to true if there are no errors getting stages', async () => { - stub(DvcReader.prototype, 'listStages').resolves('') + stub(DvcReader.prototype, 'stageList').resolves('') stub(FileSystem, 'hasDvcYamlFile').returns(false) const { experiments, messageSpy } = buildExperiments({ @@ -223,7 +223,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasConfig to false if there are no stages', async () => { - stub(DvcReader.prototype, 'listStages').resolves('') + stub(DvcReader.prototype, 'stageList').resolves('') const { experiments, messageSpy } = buildExperiments({ disposer: disposable @@ -237,7 +237,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasConfig to true if there are stages', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, messageSpy } = buildExperiments({ disposer: disposable @@ -753,7 +753,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should be able to handle a message to modify the workspace params and queue an experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, dvcExecutor } = buildExperiments({ disposer: disposable }) @@ -789,7 +789,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should be able to handle a message to modify the workspace params and run a new experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, dvcRunner } = buildExperiments({ disposer: disposable }) @@ -827,7 +827,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should be able to handle a message to modify the workspace params, reset and run a new experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, dvcRunner } = buildExperiments({ disposer: disposable }) @@ -1033,7 +1033,7 @@ suite('Experiments Test Suite', () => { }) it('should be able to handle a message to select columns', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { columnsModel, experiments, messageSpy } = setupExperimentsAndMockCommands() @@ -1398,7 +1398,7 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should handle a message to add a configuration', async () => { - stub(DvcReader.prototype, 'listStages').resolves('') + stub(DvcReader.prototype, 'stageList').resolves('') const { experiments, mockCheckOrAddPipeline, messageSpy } = setupExperimentsAndMockCommands() diff --git a/extension/src/test/suite/experiments/model/filterBy/tree.test.ts b/extension/src/test/suite/experiments/model/filterBy/tree.test.ts index 90700477d9..9ce66e9a44 100644 --- a/extension/src/test/suite/experiments/model/filterBy/tree.test.ts +++ b/extension/src/test/suite/experiments/model/filterBy/tree.test.ts @@ -61,7 +61,7 @@ suite('Experiments Filter By Tree Test Suite', () => { }) it('should be able to update the table data by adding and removing a filter', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, messageSpy } = buildExperiments({ disposer: disposable @@ -341,7 +341,7 @@ suite('Experiments Filter By Tree Test Suite', () => { }) it('should be able to filter to starred experiments', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, messageSpy } = buildExperiments({ disposer: disposable }) diff --git a/extension/src/test/suite/experiments/model/tree.test.ts b/extension/src/test/suite/experiments/model/tree.test.ts index 549868cced..f5ce6a1bf6 100644 --- a/extension/src/test/suite/experiments/model/tree.test.ts +++ b/extension/src/test/suite/experiments/model/tree.test.ts @@ -574,7 +574,7 @@ suite('Experiments Tree Test Suite', () => { }) it('should be able to queue an experiment from an existing one with dvc.views.experiments.queueExperiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { dvcExecutor, experiments, experimentsModel } = buildExperiments({ disposer: disposable @@ -635,7 +635,7 @@ suite('Experiments Tree Test Suite', () => { }) it('should be able to run a new experiment from an existing one with dvc.views.experiments.runExperiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { dvcRunner, experiments, experimentsModel } = buildExperiments({ disposer: disposable @@ -693,7 +693,7 @@ suite('Experiments Tree Test Suite', () => { }) it('should be able to reset and run a new checkpoint experiment from an existing one with dvc.views.experiments.resetAndRunCheckpointExperiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { dvcRunner, experiments, experimentsModel } = buildExperiments({ disposer: disposable diff --git a/extension/src/test/suite/experiments/workspace.test.ts b/extension/src/test/suite/experiments/workspace.test.ts index 2f61f0096d..537cebac54 100644 --- a/extension/src/test/suite/experiments/workspace.test.ts +++ b/extension/src/test/suite/experiments/workspace.test.ts @@ -111,7 +111,7 @@ suite('Workspace Experiments Test Suite', () => { DvcRunner.prototype, 'runExperiment' ).resolves(undefined) - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { workspaceExperiments, experiments } = buildMultiRepoExperiments(disposable) @@ -160,7 +160,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndQueue', () => { it('should be able to queue an experiment using an existing one as a base', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { dvcExecutor, experiments } = buildExperiments({ disposer: disposable @@ -210,7 +210,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndResume', () => { it('should be able to resume a checkpoint experiment using an existing one as a base', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments } = buildExperiments({ disposer: disposable }) @@ -262,7 +262,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndRun', () => { it('should be able to run an experiment using an existing one as a base', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments } = buildExperiments({ disposer: disposable }) @@ -314,7 +314,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.queueExperiment', () => { it('should be able to queue an experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const mockExperimentRunQueue = stub( DvcExecutor.prototype, @@ -330,7 +330,7 @@ suite('Workspace Experiments Test Suite', () => { }) it('should send a telemetry event containing a duration when an experiment is queued', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') stub(DvcExecutor.prototype, 'expRunQueue').resolves('true') @@ -352,7 +352,7 @@ suite('Workspace Experiments Test Suite', () => { }) it('should send a telemetry event containing an error message when an experiment fails to queue', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const duration = 77777 mockDuration(duration) @@ -386,7 +386,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.runExperiment', () => { it('should be able to run an experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const mockRunExperiment = stub( DvcRunner.prototype, 'runExperiment' @@ -407,7 +407,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.resumeCheckpointExperiment', () => { it('should be able to run an experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const mockRunExperiment = stub( DvcRunner.prototype, @@ -425,7 +425,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.resetAndRunCheckpointExperiment', () => { it('should be able to reset existing checkpoints and restart the experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const mockRunExperimentReset = stub( DvcRunner.prototype, 'runExperimentReset' @@ -504,7 +504,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.startExperimentsQueue', () => { it('should be able to start the experiments queue with the selected number of workers', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const mockQueueStart = stub(DvcExecutor.prototype, 'queueStart').resolves( undefined @@ -531,7 +531,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.stopExperimentsQueue', () => { it('should be able to stop the experiments queue', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const mockQueueStop = stub(DvcExecutor.prototype, 'queueStop').resolves( undefined @@ -649,7 +649,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.branchExperiment', () => { it('should be able to create a branch from an experiment', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments } = buildExperiments({ disposer: disposable }) await experiments.isReady() @@ -898,7 +898,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.removeExperimentQueue', () => { it('should remove all queued experiments from the selected repository', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments } = buildExperiments({ disposer: disposable }) diff --git a/extension/src/test/suite/extension.test.ts b/extension/src/test/suite/extension.test.ts index c32de396d6..1277c3e9f9 100644 --- a/extension/src/test/suite/extension.test.ts +++ b/extension/src/test/suite/extension.test.ts @@ -219,7 +219,7 @@ suite('Extension Test Suite', () => { describe('dvc.stopAllRunningExperiments', () => { it('should send a telemetry event containing properties relating to the event', async () => { - stub(DvcReader.prototype, 'listStages').resolves('train') + stub(DvcReader.prototype, 'stageList').resolves('train') const duration = 1234 const otherRoot = resolve('other', 'root') From b622744516d057de7e0eeeeed8ccdfbc56d69dc3 Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Wed, 12 Jul 2023 18:52:02 +1000 Subject: [PATCH 02/12] start to move functionality across to pipeline --- extension/src/data/index.ts | 15 +- extension/src/experiments/index.ts | 29 ++-- extension/src/experiments/webview/messages.ts | 30 +--- extension/src/experiments/workspace.ts | 13 +- extension/src/extension.ts | 12 +- extension/src/pipeline/data.ts | 30 +++- extension/src/pipeline/index.ts | 141 +++++++++++++++++- extension/src/pipeline/model/collect.test.ts | 56 ++++--- extension/src/pipeline/model/collect.ts | 48 +++--- extension/src/pipeline/model/index.ts | 43 ++++-- .../src/test/suite/experiments/index.test.ts | 43 +++--- extension/src/test/suite/experiments/util.ts | 30 +++- .../test/suite/experiments/workspace.test.ts | 2 +- extension/src/test/suite/pipeline/util.ts | 19 +++ extension/src/test/suite/plots/util.ts | 9 +- 15 files changed, 367 insertions(+), 153 deletions(-) create mode 100644 extension/src/test/suite/pipeline/util.ts diff --git a/extension/src/data/index.ts b/extension/src/data/index.ts index 6e3b762d33..7eef4fd4a3 100644 --- a/extension/src/data/index.ts +++ b/extension/src/data/index.ts @@ -20,10 +20,9 @@ export abstract class BaseData< T extends | { data: PlotsOutputOrError; revs: string[] } | ExperimentsOutput - | { dag: string; stageList: string } + | { dag: string; stages: { [pipeline: string]: string | undefined } } > extends DeferredDisposable { public readonly onDidUpdate: Event - public readonly onDidChangeDvcYaml: Event protected readonly dvcRoot: string protected readonly processManager: ProcessManager @@ -36,10 +35,6 @@ export abstract class BaseData< new EventEmitter() ) - private readonly dvcYamlChanged: EventEmitter = this.dispose.track( - new EventEmitter() - ) - constructor( dvcRoot: string, internalCommands: InternalCommands, @@ -56,8 +51,6 @@ export abstract class BaseData< this.onDidUpdate = this.updated.event this.staticFiles = staticFiles - this.onDidChangeDvcYaml = this.dvcYamlChanged.event - this.watchFiles() this.waitForInitialData() @@ -96,13 +89,9 @@ export abstract class BaseData< ) { void this.managedUpdate(path) } - - if (path.endsWith('dvc.yaml')) { - this.dvcYamlChanged.fire() - } } ) } - abstract managedUpdate(path?: string): Promise + abstract managedUpdate(path?: string): Promise } diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 21b60fde79..161415ecf0 100644 --- a/extension/src/experiments/index.ts +++ b/extension/src/experiments/index.ts @@ -44,6 +44,7 @@ import { ConfigKey, getConfigValue, setUserConfigValue } from '../vscode/config' import { checkSignalFile, pollSignalFileForProcess } from '../fileSystem' import { DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants' import { Response } from '../vscode/response' +import { Pipeline } from '../pipeline' export const ExperimentsScale = { ...omit(ColumnType, 'TIMESTAMP'), @@ -64,8 +65,9 @@ export class Experiments extends BaseRepository { public readonly viewKey = ViewKey.EXPERIMENTS - private readonly data: ExperimentsData + private readonly pipeline: Pipeline + private readonly data: ExperimentsData private readonly experiments: ExperimentsModel private readonly columns: ColumnsModel @@ -96,7 +98,6 @@ export class Experiments extends BaseRepository { private dvcLiveOnlyCleanupInitialized = false private dvcLiveOnlySignalFile: string - private readonly addStage: () => Promise private readonly selectBranches: ( branchesSelected: string[] ) => Promise @@ -104,9 +105,9 @@ export class Experiments extends BaseRepository { constructor( dvcRoot: string, internalCommands: InternalCommands, + pipeline: Pipeline, resourceLocator: ResourceLocator, workspaceState: Memento, - addStage: () => Promise, selectBranches: ( branchesSelected: string[] ) => Promise, @@ -120,7 +121,7 @@ export class Experiments extends BaseRepository { ) this.internalCommands = internalCommands - this.addStage = addStage + this.pipeline = pipeline this.selectBranches = selectBranches this.onDidChangeIsExperimentsFileFocused = this.experimentsFileFocused.event @@ -146,11 +147,6 @@ export class Experiments extends BaseRepository { ) this.dispose.track(this.data.onDidUpdate(data => this.setState(data))) - this.dispose.track( - this.data.onDidChangeDvcYaml(() => - this.webviewMessages.changeHasConfig(true) - ) - ) this.dispose.track( workspace.onDidChangeConfiguration((event: ConfigurationChangeEvent) => { @@ -527,10 +523,16 @@ export class Experiments extends BaseRepository { private setupInitialData() { const waitForInitialData = this.dispose.track( - this.onDidChangeExperiments(() => { + this.onDidChangeExperiments(async () => { + await this.pipeline.isReady() this.deferred.resolve() this.dispose.untrack(waitForInitialData) waitForInitialData.dispose() + this.dispose.track( + this.pipeline.onDidUpdate(() => + this.webviewMessages.sendWebviewMessage() + ) + ) }) ) } @@ -555,15 +557,10 @@ export class Experiments extends BaseRepository { this.dvcRoot, this.experiments, this.columns, + this.pipeline, () => this.getWebview(), () => this.notifyChanged(), () => this.selectColumns(), - () => - this.internalCommands.executeCommand( - AvailableCommands.STAGE_LIST, - this.dvcRoot - ), - () => this.addStage(), (branchesSelected: string[]) => this.selectBranches(branchesSelected), () => this.data.update() ) diff --git a/extension/src/experiments/webview/messages.ts b/extension/src/experiments/webview/messages.ts index 103c5730c5..3985bcdca9 100644 --- a/extension/src/experiments/webview/messages.ts +++ b/extension/src/experiments/webview/messages.ts @@ -20,25 +20,20 @@ import { SortDefinition } from '../model/sortBy' import { getPositiveIntegerInput } from '../../vscode/inputBox' import { Title } from '../../vscode/title' import { ConfigKey, setConfigValue } from '../../vscode/config' -import { hasDvcYamlFile } from '../../fileSystem' import { NUM_OF_COMMITS_TO_INCREASE } from '../../cli/dvc/constants' +import { Pipeline } from '../../pipeline' export class WebviewMessages { private readonly dvcRoot: string private readonly experiments: ExperimentsModel private readonly columns: ColumnsModel + private readonly pipeline: Pipeline private readonly getWebview: () => BaseWebview | undefined private readonly notifyChanged: () => void private readonly selectColumns: () => Promise - private readonly hasStages: () => Promise - - private hasConfig = false - private hasValidDvcYaml = true - - private readonly addStage: () => Promise private readonly selectBranches: ( branchesSelected: string[] ) => Promise @@ -49,11 +44,10 @@ export class WebviewMessages { dvcRoot: string, experiments: ExperimentsModel, columns: ColumnsModel, + pipeline: Pipeline, getWebview: () => BaseWebview | undefined, notifyChanged: () => void, selectColumns: () => Promise, - hasStages: () => Promise, - addStage: () => Promise, selectBranches: ( branchesSelected: string[] ) => Promise, @@ -62,22 +56,12 @@ export class WebviewMessages { this.dvcRoot = dvcRoot this.experiments = experiments this.columns = columns + this.pipeline = pipeline this.getWebview = getWebview this.notifyChanged = notifyChanged this.selectColumns = selectColumns - this.hasStages = hasStages - this.addStage = addStage this.selectBranches = selectBranches this.update = update - - void this.changeHasConfig() - } - - public async changeHasConfig(update?: boolean) { - const stages = await this.hasStages() - this.hasValidDvcYaml = !hasDvcYamlFile(this.dvcRoot) || stages !== undefined - this.hasConfig = !!stages - update && this.sendWebviewMessage() } public sendWebviewMessage() { @@ -267,11 +251,11 @@ export class WebviewMessages { this.experiments.getAvailableBranchesToShow().length > 0, hasCheckpoints: this.experiments.hasCheckpoints(), hasColumns: this.columns.hasNonDefaultColumns(), - hasConfig: this.hasConfig, + hasConfig: !!this.pipeline.hasStage(), hasMoreCommits: this.experiments.getHasMoreCommits(), hasRunningWorkspaceExperiment: this.experiments.hasRunningWorkspaceExperiment(), - hasValidDvcYaml: this.hasValidDvcYaml, + hasValidDvcYaml: !this.pipeline.hasInvalidRootDvcYaml(), isShowingMoreCommits: this.experiments.getIsShowingMoreCommits(), rows: this.experiments.getRowData(), selectedForPlotsCount: this.experiments.getSelectedRevisions().length, @@ -280,7 +264,7 @@ export class WebviewMessages { } private async addConfiguration() { - await this.addStage() + await this.pipeline.checkOrAddPipeline() } private async setMaxTableHeadDepth() { diff --git a/extension/src/experiments/workspace.ts b/extension/src/experiments/workspace.ts index 32bbdbe196..8f7d5bf511 100644 --- a/extension/src/experiments/workspace.ts +++ b/extension/src/experiments/workspace.ts @@ -31,6 +31,7 @@ import { } from '../fileSystem' import { quickPickManyValues, quickPickOneOrInput } from '../vscode/quickPick' import { pickFile } from '../vscode/resourcePicker' +import { WorkspacePipeline } from '../pipeline/workspace' export enum scriptCommand { JUPYTER = 'jupyter nbconvert --to notebook --inplace --execute', @@ -201,7 +202,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< if (!cwd) { return } - const shouldContinue = await this.checkOrAddPipeline(cwd) + const shouldContinue = await this.checkOrAddPipeline(cwd) // repository.getPipeLineCwd() -> !cwd return if (!shouldContinue) { return } @@ -306,14 +307,18 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< ) } - public createRepository(dvcRoot: string, resourceLocator: ResourceLocator) { + public createRepository( + dvcRoot: string, + pipeline: WorkspacePipeline, + resourceLocator: ResourceLocator + ) { const experiments = this.dispose.track( new Experiments( dvcRoot, this.internalCommands, + pipeline.getRepository(dvcRoot), resourceLocator, this.workspaceState, - () => this.checkOrAddPipeline(dvcRoot), (branchesSelected: string[]) => this.selectBranches(branchesSelected) ) ) @@ -451,7 +456,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< if (!cwd) { return } - const shouldContinue = await this.checkOrAddPipeline(cwd) + const shouldContinue = await this.checkOrAddPipeline(cwd) // repository.getPipeLineCwd() -> !cwd return if (!shouldContinue) { return } diff --git a/extension/src/extension.ts b/extension/src/extension.ts index d3d4a55f82..694a19e21e 100644 --- a/extension/src/extension.ts +++ b/extension/src/extension.ts @@ -280,15 +280,20 @@ class Extension extends Disposable { await Promise.all([ this.repositories.create(this.getRoots()), this.repositoriesTree.initialize(this.getRoots()), - this.experiments.create(this.getRoots(), this.resourceLocator), this.pipelines.create(this.getRoots()) ]) + this.experiments.create( + this.getRoots(), + this.pipelines, + this.resourceLocator + ) this.plots.create(this.getRoots(), this.resourceLocator, this.experiments) return Promise.all([ - this.repositories.isReady(), this.experiments.isReady(), - this.plots.isReady() + this.pipelines.isReady(), + this.plots.isReady(), + this.repositories.isReady() ]) } @@ -296,6 +301,7 @@ class Extension extends Disposable { this.repositories.reset() this.repositoriesTree.initialize([]) this.experiments.reset() + this.pipelines.reset() this.plots.reset() } diff --git a/extension/src/pipeline/data.ts b/extension/src/pipeline/data.ts index bf6dd09c6c..b7c33cc0a2 100644 --- a/extension/src/pipeline/data.ts +++ b/extension/src/pipeline/data.ts @@ -1,7 +1,12 @@ +import { dirname } from 'path' import { AvailableCommands, InternalCommands } from '../commands/internal' import { BaseData } from '../data' +import { findFiles } from '../fileSystem/workspace' -export class PipelineData extends BaseData<{ dag: string; stageList: string }> { +export class PipelineData extends BaseData<{ + dag: string + stages: { [pipeline: string]: string | undefined } +}> { constructor(dvcRoot: string, internalCommands: InternalCommands) { super( dvcRoot, @@ -18,13 +23,26 @@ export class PipelineData extends BaseData<{ dag: string; stageList: string }> { } public async update(): Promise { - const [dag, stageList] = await Promise.all([ + const [dag, fileList] = await Promise.all([ this.internalCommands.executeCommand(AvailableCommands.DAG, this.dvcRoot), - this.internalCommands.executeCommand( + findFiles('**/dvc.yaml') + ]) + + const dvcYamlsDirs = new Set() + for (const file of fileList) { + if (file.startsWith(this.dvcRoot)) { + dvcYamlsDirs.add(dirname(file)) + } + } + + const stages: { [dir: string]: string } = {} + for (const dir of dvcYamlsDirs) { + stages[dir] = await this.internalCommands.executeCommand( AvailableCommands.STAGE_LIST, - this.dvcRoot + dir ) - ]) - return this.notifyChanged({ dag, stageList }) + } + + return this.notifyChanged({ dag, stages }) } } diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index 6d5bad767b..aaec8e09f1 100644 --- a/extension/src/pipeline/index.ts +++ b/extension/src/pipeline/index.ts @@ -1,12 +1,39 @@ import { join } from 'path' +import { Event, EventEmitter } from 'vscode' import { appendFileSync, writeFileSync } from 'fs-extra' import { PipelineData } from './data' import { PipelineModel } from './model' import { DeferredDisposable } from '../class/deferred' import { InternalCommands } from '../commands/internal' import { TEMP_DAG_FILE } from '../cli/dvc/constants' +import { findOrCreateDvcYamlFile, getFileExtension } from '../fileSystem' +import { Toast } from '../vscode/toast' +import { getInput, getValidInput } from '../vscode/inputBox' +import { Title } from '../vscode/title' +import { quickPickOneOrInput } from '../vscode/quickPick' +import { pickFile } from '../vscode/resourcePicker' + +enum scriptCommand { + JUPYTER = 'jupyter nbconvert --to notebook --inplace --execute', + PYTHON = 'python' +} + +const getScriptCommand = (script: string) => { + switch (getFileExtension(script)) { + case '.py': + return scriptCommand.PYTHON + case '.ipynb': + return scriptCommand.JUPYTER + default: + return '' + } +} export class Pipeline extends DeferredDisposable { + public onDidUpdate: Event + + private updated: EventEmitter + private readonly dvcRoot: string private readonly data: PipelineData private readonly model: PipelineModel @@ -16,23 +43,131 @@ export class Pipeline extends DeferredDisposable { this.dvcRoot = dvcRoot this.data = this.dispose.track(new PipelineData(dvcRoot, internalCommands)) this.model = this.dispose.track(new PipelineModel(dvcRoot)) + this.updated = this.dispose.track(new EventEmitter()) + this.onDidUpdate = this.updated.event void this.initialize() } + public hasStage() { + return this.model.hasStage() + } + + public hasInvalidRootDvcYaml() { + return this.model.hasInvalidRootDvcYaml() + } + + public async checkOrAddPipeline() { + const hasStage = this.model.hasStage() + if (this.hasInvalidRootDvcYaml()) { + await Toast.showError( + 'Cannot perform task. Your dvc.yaml file is invalid.' + ) + return false + } + + if (!hasStage) { + return this.addPipeline() + } + return true + } + public forceRerender() { return appendFileSync(join(this.dvcRoot, TEMP_DAG_FILE), '\n') } private async initialize() { this.dispose.track( - this.data.onDidUpdate(({ dag, stageList }) => { - writeFileSync(join(this.dvcRoot, TEMP_DAG_FILE), dag) - this.model.transformAndSet(stageList) + this.data.onDidUpdate(({ dag, stages }) => { + this.writeDag(dag) + const hasStage = this.model.hasStage() + const hasInvalidRootDvcYaml = this.model.hasInvalidRootDvcYaml() + this.model.transformAndSet(stages) + if ( + hasStage !== this.model.hasStage() || + hasInvalidRootDvcYaml !== this.model.hasInvalidRootDvcYaml() + ) { + this.updated.fire() + } }) ) await this.data.isReady() return this.deferred.resolve() } + + private async addPipeline() { + const stageName = await this.askForStageName() + if (!stageName) { + return false + } + + const { trainingScript, command, enteredManually } = + await this.askForTrainingScript() + if (!trainingScript) { + return false + } + void findOrCreateDvcYamlFile( + this.dvcRoot, + trainingScript, + stageName, + command, + !enteredManually + ) + return true + } + + private async askForStageName() { + return await getValidInput( + Title.ENTER_STAGE_NAME, + (stageName?: string) => { + if (!stageName) { + return 'Stage name must not be empty' + } + if (!/^[a-z]/i.test(stageName)) { + return 'Stage name should start with a letter' + } + return /^\w+$/.test(stageName) + ? null + : 'Stage name should only include letters and numbers' + }, + { value: 'train' } + ) + } + + private async askForTrainingScript() { + const selectValue = 'select' + const pathOrSelect = await quickPickOneOrInput( + [{ label: 'Select from file explorer', value: selectValue }], + { + defaultValue: '', + placeholder: 'Path to script', + title: Title.ENTER_PATH_OR_CHOOSE_FILE + } + ) + + const trainingScript = + pathOrSelect === selectValue + ? await pickFile(Title.SELECT_TRAINING_SCRIPT) + : pathOrSelect + + if (!trainingScript) { + return { + command: undefined, + enteredManually: false, + trainingScript: undefined + } + } + + const command = + getScriptCommand(trainingScript) || + (await getInput(Title.ENTER_COMMAND_TO_RUN)) || + '' + const enteredManually = pathOrSelect !== selectValue + return { command, enteredManually, trainingScript } + } + + private writeDag(dag: string) { + writeFileSync(join(this.dvcRoot, TEMP_DAG_FILE), dag) + } } diff --git a/extension/src/pipeline/model/collect.test.ts b/extension/src/pipeline/model/collect.test.ts index 290c893ab7..6bb096a34b 100644 --- a/extension/src/pipeline/model/collect.test.ts +++ b/extension/src/pipeline/model/collect.test.ts @@ -1,44 +1,38 @@ import { join } from 'path' -import { collectPipelines, collectStages } from './collect' +import { collectStages } from './collect' import { dvcDemoPath } from '../../test/util' +// need to combine collect stages and collect pipelines as a pipeline is only valid if it has stages + describe('collectStages', () => { it('should handle a simple list of stages', () => { - const stages = collectStages( - 'data.dvc Outputs data\ntrain Outputs model.pt, training/plots, hist.csv; Reports training/metrics.json' - ) - expect(stages).toStrictEqual(['train']) + const { validPipelines, validStages } = collectStages({ + [dvcDemoPath]: + 'data.dvc Outputs data\ntrain Outputs model.pt, training/plots, hist.csv; Reports training/metrics.json' + }) + expect(validStages).toStrictEqual(['train']) + expect(validPipelines).toStrictEqual(new Set([dvcDemoPath])) }) it('should handle a monorepos list of stages', () => { - const stages = collectStages( - `nested1/data/data.xml.dvc Outputs data.xml - nested1/dvc.yaml:prepare Outputs data/prepared - nested1/dvc.yaml:featurize Outputs data/features - nested1/dvc.yaml:train Outputs model.pkl - nested1/dvc.yaml:evaluate Outputs eval/importance.png, eval/live/plots, eval/prc; Reports eval/live/metri… - nested2/data/data.xml.dvc Outputs data.xml` - ) - expect(stages).toStrictEqual([ - 'nested1/dvc.yaml:prepare', - 'nested1/dvc.yaml:featurize', - 'nested1/dvc.yaml:train', - 'nested1/dvc.yaml:evaluate' - ]) - }) -}) - -describe('collectPipelines', () => { - it('should handle a list of stages', () => { - const pipelines = collectPipelines(dvcDemoPath, [ + const { validStages, validPipelines, invalidPipelines } = collectStages({ + [join(dvcDemoPath, 'nested1')]: `data/data.xml.dvc Outputs data.xml + prepare Outputs data/prepared + featurize Outputs data/features + train Outputs model.pkl + evaluate Outputs eval/importance.png, eval/live/plots, eval/prc; Reports eval/live/metri…`, + [join('dvcDemoPath', 'nested2')]: 'data/data.xml.dvc Outputs data.xml' + }) + expect(validStages).toStrictEqual([ + 'prepare', + 'featurize', 'train', - 'nested1/dvc.yaml:prepare', - 'nested1/dvc.yaml:featurize', - 'nested1/dvc.yaml:train', - 'nested1/dvc.yaml:evaluate' + 'evaluate' ]) - expect(pipelines).toStrictEqual( - new Set([dvcDemoPath, join(dvcDemoPath, 'nested1')]) + + expect(validPipelines).toStrictEqual( + new Set([join(dvcDemoPath, 'nested1')]) ) + expect(invalidPipelines).toStrictEqual(new Set()) }) }) diff --git a/extension/src/pipeline/model/collect.ts b/extension/src/pipeline/model/collect.ts index b0253d183a..e561053724 100644 --- a/extension/src/pipeline/model/collect.ts +++ b/extension/src/pipeline/model/collect.ts @@ -1,31 +1,33 @@ -import { join } from 'path' import { trimAndSplit } from '../../util/stdout' -export const collectStages = (stageList: string): string[] => { - const stages: string[] = [] - for (const stageStr of trimAndSplit(stageList)) { - const stage = stageStr.split(/\s+/)?.[0]?.trim() - if (!stage || stage.endsWith('.dvc')) { - continue +export const collectStages = (pipelines: { + [pipeline: string]: string | undefined +}): { + invalidPipelines: Set + validPipelines: Set + validStages: string[] + // eslint-disable-next-line sonarjs/cognitive-complexity +} => { + const validStages: string[] = [] + const validPipelines = new Set() + const invalidPipelines = new Set() + + for (const [pipeline, stageList] of Object.entries(pipelines)) { + if (stageList === undefined) { + invalidPipelines.add(pipeline) } - stages.push(stage) - } - return stages -} -export const collectPipelines = ( - dvcRoot: string, - stages: string[] -): Set => { - const dvcYamlIdentifier = '/dvc.yaml:' - const pipelines = new Set() - for (const stage of stages) { - if (!stage.includes(dvcYamlIdentifier)) { - pipelines.add(dvcRoot) + if (!stageList) { continue } - const [pipeline] = stage.split(dvcYamlIdentifier) - pipelines.add(join(dvcRoot, ...pipeline.split('/'))) + for (const stageStr of trimAndSplit(stageList)) { + const stage = stageStr.split(/\s+/)?.[0]?.trim() + if (!stage || stage.endsWith('.dvc')) { + continue + } + validStages.push(stage) + validPipelines.add(pipeline) + } } - return pipelines + return { invalidPipelines, validPipelines, validStages } } diff --git a/extension/src/pipeline/model/index.ts b/extension/src/pipeline/model/index.ts index 205db14ddb..cbd567832c 100644 --- a/extension/src/pipeline/model/index.ts +++ b/extension/src/pipeline/model/index.ts @@ -1,26 +1,51 @@ -import { collectPipelines, collectStages } from './collect' +import isEqual from 'lodash.isequal' +import { collectStages } from './collect' import { Disposable } from '../../class/dispose' export class PipelineModel extends Disposable { private readonly dvcRoot: string - private stages: string[] = [] - private pipelines: Set = new Set() + private stages: string[] | undefined = [] + private validPipelines: Set | undefined + private invalidPipelines: Set | undefined constructor(dvcRoot: string) { super() this.dvcRoot = dvcRoot } - public transformAndSet(stageList: string) { - this.stages = collectStages(stageList) - this.pipelines = collectPipelines(this.dvcRoot, this.stages) + public transformAndSet(stages: { [pipeline: string]: string | undefined }) { + if (isEqual(stages, {})) { + this.stages = undefined + this.validPipelines = undefined + this.invalidPipelines = undefined + return + } + + const { validPipelines, invalidPipelines, validStages } = + collectStages(stages) + + this.validPipelines = validPipelines + this.invalidPipelines = invalidPipelines + this.stages = validStages + } + + public hasStage() { + return !!(this.stages && this.stages.length > 0) } - public hasStages() { - return this.stages.length > 0 + public getStages() { + return this.stages } public getPipelines() { - return this.pipelines + return this.validPipelines + } + + public hasPipeline() { + return !!(this.validPipelines && this.validPipelines.size > 0) + } + + public hasInvalidRootDvcYaml() { + return !!this.invalidPipelines?.has(this.dvcRoot) } } diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 39e4e93970..5cd12ed4c8 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -91,6 +91,8 @@ import { DEFAULT_NB_ITEMS_PER_ROW } from '../../../plots/webview/contract' import { Toast } from '../../../vscode/toast' import { Response } from '../../../vscode/response' import { MAX_SELECTED_EXPERIMENTS } from '../../../experiments/model/status' +import { Pipeline } from '../../../pipeline' +import { buildPipeline } from '../pipeline/util' const { openFileInEditor } = FileSystem @@ -175,20 +177,24 @@ suite('Experiments Test Suite', () => { expect(windowSpy).not.to.have.been.called }).timeout(WEBVIEW_TEST_TIMEOUT) - it('should set hasValidDvcYaml to false if there is an error getting stages and there is a dvc.yaml file', async () => { - stub(DvcReader.prototype, 'stageList').resolves(undefined) - stub(FileSystem, 'hasDvcYamlFile').returns(true) + it.only( + 'should set hasValidDvcYaml to false if there is an error getting stages and there is a dvc.yaml file', + async () => { + stub(FileSystem, 'hasDvcYamlFile').returns(true) - const { experiments, messageSpy } = buildExperiments({ - disposer: disposable - }) + const { experiments, messageSpy, dvcReader } = buildExperiments({ + disposer: disposable + }) - await experiments.showWebview() + stub(dvcReader, 'stageList').resolves(undefined) - expect(messageSpy).to.be.calledWithMatch({ - hasValidDvcYaml: false - }) - }).timeout(WEBVIEW_TEST_TIMEOUT) + await experiments.showWebview() + + expect(messageSpy).to.be.calledWithMatch({ + hasValidDvcYaml: false + }) + } + ).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasValidDvcYaml to true if there is an error getting stages and there is no dvc.yaml file', async () => { stub(DvcReader.prototype, 'stageList').resolves(undefined) @@ -236,9 +242,7 @@ suite('Experiments Test Suite', () => { }) }).timeout(WEBVIEW_TEST_TIMEOUT) - it('should set hasConfig to true if there are stages', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - + it.only('should set hasConfig to true if there are stages', async () => { const { experiments, messageSpy } = buildExperiments({ disposer: disposable }) @@ -1563,14 +1567,19 @@ suite('Experiments Test Suite', () => { const resourceLocator = disposable.track( new ResourceLocator(extensionUri) ) + const pipeline = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + internalCommands + }) const experiments = disposable.track( new Experiments( dvcDemoPath, internalCommands, + pipeline, resourceLocator, buildMockMemento(), - () => Promise.resolve(true), () => Promise.resolve([]), buildMockExperimentsData() ) @@ -1743,9 +1752,9 @@ suite('Experiments Test Suite', () => { new Experiments( 'test', internalCommands, + { hasStage: () => true } as Pipeline, {} as ResourceLocator, mockMemento, - () => Promise.resolve(true), () => Promise.resolve([]), buildMockExperimentsData() ) @@ -1904,9 +1913,9 @@ suite('Experiments Test Suite', () => { new Experiments( 'test', internalCommands, + { hasStage: () => true } as Pipeline, {} as ResourceLocator, mockMemento, - () => Promise.resolve(true), () => Promise.resolve([]), buildMockExperimentsData() ) diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index a6ea7d7ecd..ae09aee4df 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -19,6 +19,7 @@ import { ColumnsModel } from '../../../experiments/columns/model' import { DEFAULT_NUM_OF_COMMITS_TO_SHOW } from '../../../cli/dvc/constants' import { PersistenceKey } from '../../../persistence/constants' import { ExpShowOutput } from '../../../cli/dvc/contract' +import { buildPipeline } from '../pipeline/util' export const DEFAULT_EXPERIMENTS_OUTPUT = { availableNbCommits: { main: 5 }, @@ -33,7 +34,8 @@ export const buildExperiments = ({ dvcRoot = dvcDemoPath, expShow = expShowFixture, gitLog = gitLogFixture, - rowOrder = rowOrderFixture + rowOrder = rowOrderFixture, + stageList = 'train' }: { availableNbCommits?: { [branch: string]: number } disposer: Disposer @@ -41,6 +43,7 @@ export const buildExperiments = ({ expShow?: ExpShowOutput gitLog?: string rowOrder?: { branch: string; sha: string }[] + stageList?: string | null }) => { const { dvcExecutor, @@ -60,7 +63,12 @@ export const buildExperiments = ({ const mockExperimentsData = buildMockExperimentsData( mockUpdateExperimentsData ) - const mockCheckOrAddPipeline = stub() + + stub(dvcReader, 'stageList').resolves(stageList ?? undefined) + stub(dvcReader, 'dag').resolves('') + + const pipeline = buildPipeline({ disposer, dvcRoot, internalCommands }) + const mockCheckOrAddPipeline = stub(pipeline, 'checkOrAddPipeline') const mockSelectBranches = stub().resolves(['main', 'other']) const mockMemento = buildMockMemento({ [`${PersistenceKey.EXPERIMENTS_BRANCHES}${dvcRoot}`]: ['main'], @@ -73,9 +81,9 @@ export const buildExperiments = ({ new Experiments( dvcRoot, internalCommands, + pipeline, resourceLocator, mockMemento, - mockCheckOrAddPipeline, mockSelectBranches, mockExperimentsData ) @@ -131,8 +139,17 @@ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { 'other/dvc/root': mockExperiments }) ) + + const pipeline = buildPipeline({ + disposer, + dvcRoot: dvcDemoPath, + internalCommands + }) + stub(pipeline, 'hasStage').returns(true) + const [experiments] = workspaceExperiments.create( [dvcDemoPath], + { getRepository: () => pipeline }, resourceLocator ) @@ -148,8 +165,15 @@ export const buildSingleRepoExperiments = (disposer: SafeWatcherDisposer) => { const workspaceExperiments = disposer.track( new WorkspaceExperiments(internalCommands, buildMockMemento()) ) + const pipeline = buildPipeline({ + disposer, + dvcRoot: dvcDemoPath, + internalCommands + }) + const [experiments] = workspaceExperiments.create( [dvcDemoPath], + { getRepository: () => pipeline }, resourceLocator ) diff --git a/extension/src/test/suite/experiments/workspace.test.ts b/extension/src/test/suite/experiments/workspace.test.ts index 537cebac54..8540ff4675 100644 --- a/extension/src/test/suite/experiments/workspace.test.ts +++ b/extension/src/test/suite/experiments/workspace.test.ts @@ -90,7 +90,7 @@ suite('Workspace Experiments Test Suite', () => { expect(mockQuickPickOne).to.be.calledOnce }).timeout(WEBVIEW_TEST_TIMEOUT) - it('should not prompt to pick a project if there is only one project', async () => { + it.only('should not prompt to pick a project if there is only one project', async () => { const mockQuickPickOne = stub(QuickPick, 'quickPickOne').resolves( dvcDemoPath ) diff --git a/extension/src/test/suite/pipeline/util.ts b/extension/src/test/suite/pipeline/util.ts new file mode 100644 index 0000000000..e395de7062 --- /dev/null +++ b/extension/src/test/suite/pipeline/util.ts @@ -0,0 +1,19 @@ +import { stub } from 'sinon' +import { InternalCommands } from '../../../commands/internal' +import { Disposer } from '../../../extension' +import { Pipeline } from '../../../pipeline' + +export const buildPipeline = ({ + disposer, + dvcRoot, + internalCommands +}: { + disposer: Disposer + dvcRoot: string + internalCommands: InternalCommands +}): Pipeline => { + const pipeline = disposer.track(new Pipeline(dvcRoot, internalCommands)) + // eslint-disable-next-line @typescript-eslint/no-explicit-any + stub(pipeline as any, 'writeDag').resolves() + return pipeline +} diff --git a/extension/src/test/suite/plots/util.ts b/extension/src/test/suite/plots/util.ts index fb55b255d1..8a21b64bbd 100644 --- a/extension/src/test/suite/plots/util.ts +++ b/extension/src/test/suite/plots/util.ts @@ -26,6 +26,7 @@ import { } from '../../../cli/dvc/contract' import { ErrorsModel } from '../../../plots/errors/model' import { PersistenceKey } from '../../../persistence/constants' +import { buildPipeline } from '../pipeline/util' export const buildPlots = async ({ availableNbCommits = { main: 5 }, @@ -50,10 +51,17 @@ export const buildPlots = async ({ MOCK_IMAGE_MTIME ) + const pipeline = buildPipeline({ + disposer, + dvcRoot: dvcDemoPath, + internalCommands + }) + const experiments = disposer.track( new Experiments( dvcDemoPath, internalCommands, + pipeline, resourceLocator, buildMockMemento({ [`${PersistenceKey.EXPERIMENTS_BRANCHES}${dvcDemoPath}`]: ['main'], @@ -61,7 +69,6 @@ export const buildPlots = async ({ main: 5 } }), - () => Promise.resolve(true), () => Promise.resolve([]), buildMockExperimentsData() ) From e80d0f6de75e502c5de2e586267f14256d631b45 Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Thu, 13 Jul 2023 10:57:04 +1000 Subject: [PATCH 03/12] fix integration tests --- .../src/test/suite/experiments/index.test.ts | 69 ++++++++----------- .../experiments/model/filterBy/tree.test.ts | 4 -- .../test/suite/experiments/model/tree.test.ts | 7 -- extension/src/test/suite/experiments/util.ts | 1 + .../test/suite/experiments/workspace.test.ts | 37 +--------- 5 files changed, 33 insertions(+), 85 deletions(-) diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 5cd12ed4c8..f1edc6866a 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -85,7 +85,6 @@ import { AvailableCommands } from '../../../commands/internal' import { Setup } from '../../../setup' import * as FileSystem from '../../../fileSystem' import * as ProcessExecution from '../../../process/execution' -import { DvcReader } from '../../../cli/dvc/reader' import { DvcViewer } from '../../../cli/dvc/viewer' import { DEFAULT_NB_ITEMS_PER_ROW } from '../../../plots/webview/contract' import { Toast } from '../../../vscode/toast' @@ -124,8 +123,6 @@ suite('Experiments Test Suite', () => { describe('showWebview', () => { it('should be able to make the experiment webview visible', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { experiments, messageSpy } = buildExperiments({ disposer: disposable }) @@ -177,33 +174,27 @@ suite('Experiments Test Suite', () => { expect(windowSpy).not.to.have.been.called }).timeout(WEBVIEW_TEST_TIMEOUT) - it.only( - 'should set hasValidDvcYaml to false if there is an error getting stages and there is a dvc.yaml file', - async () => { - stub(FileSystem, 'hasDvcYamlFile').returns(true) - - const { experiments, messageSpy, dvcReader } = buildExperiments({ - disposer: disposable - }) - - stub(dvcReader, 'stageList').resolves(undefined) - - await experiments.showWebview() + it('should set hasValidDvcYaml to false if there is an error getting stages and there is a dvc.yaml file', async () => { + const { experiments, messageSpy } = buildExperiments({ + disposer: disposable, + stageList: null + }) - expect(messageSpy).to.be.calledWithMatch({ - hasValidDvcYaml: false - }) - } - ).timeout(WEBVIEW_TEST_TIMEOUT) + await experiments.showWebview() - it('should set hasValidDvcYaml to true if there is an error getting stages and there is no dvc.yaml file', async () => { - stub(DvcReader.prototype, 'stageList').resolves(undefined) - stub(FileSystem, 'hasDvcYamlFile').returns(false) + expect(messageSpy).to.be.calledWithMatch({ + hasValidDvcYaml: false + }) + }).timeout(WEBVIEW_TEST_TIMEOUT) - const { experiments, messageSpy } = buildExperiments({ - disposer: disposable + it('should set hasValidDvcYaml to true if there is no invalid root dvc.yaml file', async () => { + const { experiments, messageSpy, pipeline } = buildExperiments({ + disposer: disposable, + stageList: null }) + stub(pipeline, 'hasInvalidRootDvcYaml').returns(false) + await experiments.showWebview() const expectedTableData = { @@ -214,11 +205,11 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasValidDvcYaml to true if there are no errors getting stages', async () => { - stub(DvcReader.prototype, 'stageList').resolves('') stub(FileSystem, 'hasDvcYamlFile').returns(false) const { experiments, messageSpy } = buildExperiments({ - disposer: disposable + disposer: disposable, + stageList: '' }) await experiments.showWebview() @@ -229,10 +220,9 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should set hasConfig to false if there are no stages', async () => { - stub(DvcReader.prototype, 'stageList').resolves('') - const { experiments, messageSpy } = buildExperiments({ - disposer: disposable + disposer: disposable, + stageList: '' }) await experiments.showWebview() @@ -242,7 +232,7 @@ suite('Experiments Test Suite', () => { }) }).timeout(WEBVIEW_TEST_TIMEOUT) - it.only('should set hasConfig to true if there are stages', async () => { + it('should set hasConfig to true if there are stages', async () => { const { experiments, messageSpy } = buildExperiments({ disposer: disposable }) @@ -757,7 +747,6 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should be able to handle a message to modify the workspace params and queue an experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, dvcExecutor } = buildExperiments({ disposer: disposable }) @@ -793,7 +782,6 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should be able to handle a message to modify the workspace params and run a new experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, dvcRunner } = buildExperiments({ disposer: disposable }) @@ -831,7 +819,6 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should be able to handle a message to modify the workspace params, reset and run a new experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, dvcRunner } = buildExperiments({ disposer: disposable }) @@ -1037,8 +1024,6 @@ suite('Experiments Test Suite', () => { }) it('should be able to handle a message to select columns', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { columnsModel, experiments, messageSpy } = setupExperimentsAndMockCommands() @@ -1402,8 +1387,6 @@ suite('Experiments Test Suite', () => { }).timeout(WEBVIEW_TEST_TIMEOUT) it('should handle a message to add a configuration', async () => { - stub(DvcReader.prototype, 'stageList').resolves('') - const { experiments, mockCheckOrAddPipeline, messageSpy } = setupExperimentsAndMockCommands() @@ -1752,7 +1735,10 @@ suite('Experiments Test Suite', () => { new Experiments( 'test', internalCommands, - { hasStage: () => true } as Pipeline, + { + hasStage: () => true, + isReady: () => Promise.resolve() + } as Pipeline, {} as ResourceLocator, mockMemento, () => Promise.resolve([]), @@ -1913,7 +1899,10 @@ suite('Experiments Test Suite', () => { new Experiments( 'test', internalCommands, - { hasStage: () => true } as Pipeline, + { + hasStage: () => true, + isReady: () => Promise.resolve() + } as Pipeline, {} as ResourceLocator, mockMemento, () => Promise.resolve([]), diff --git a/extension/src/test/suite/experiments/model/filterBy/tree.test.ts b/extension/src/test/suite/experiments/model/filterBy/tree.test.ts index 9ce66e9a44..975741f965 100644 --- a/extension/src/test/suite/experiments/model/filterBy/tree.test.ts +++ b/extension/src/test/suite/experiments/model/filterBy/tree.test.ts @@ -34,7 +34,6 @@ import { FilterItem } from '../../../../../experiments/model/filterBy/tree' import { starredFilter } from '../../../../../experiments/model/filterBy/constants' -import { DvcReader } from '../../../../../cli/dvc/reader' import { Value, ValueTree, @@ -61,8 +60,6 @@ suite('Experiments Filter By Tree Test Suite', () => { }) it('should be able to update the table data by adding and removing a filter', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { experiments, messageSpy } = buildExperiments({ disposer: disposable }) @@ -341,7 +338,6 @@ suite('Experiments Filter By Tree Test Suite', () => { }) it('should be able to filter to starred experiments', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') const { experiments, messageSpy } = buildExperiments({ disposer: disposable }) diff --git a/extension/src/test/suite/experiments/model/tree.test.ts b/extension/src/test/suite/experiments/model/tree.test.ts index f5ce6a1bf6..25149e2c89 100644 --- a/extension/src/test/suite/experiments/model/tree.test.ts +++ b/extension/src/test/suite/experiments/model/tree.test.ts @@ -36,7 +36,6 @@ import { DvcExecutor } from '../../../../cli/dvc/executor' import { Param } from '../../../../experiments/model/modify/collect' import { WorkspaceExperiments } from '../../../../experiments/workspace' import { EXPERIMENT_WORKSPACE_ID } from '../../../../cli/dvc/contract' -import { DvcReader } from '../../../../cli/dvc/reader' import { copyOriginalColors } from '../../../../experiments/model/status/colors' import { Revision } from '../../../../plots/webview/contract' import { BaseWebview } from '../../../../webview' @@ -574,8 +573,6 @@ suite('Experiments Tree Test Suite', () => { }) it('should be able to queue an experiment from an existing one with dvc.views.experiments.queueExperiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { dvcExecutor, experiments, experimentsModel } = buildExperiments({ disposer: disposable }) @@ -635,8 +632,6 @@ suite('Experiments Tree Test Suite', () => { }) it('should be able to run a new experiment from an existing one with dvc.views.experiments.runExperiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { dvcRunner, experiments, experimentsModel } = buildExperiments({ disposer: disposable }) @@ -693,8 +688,6 @@ suite('Experiments Tree Test Suite', () => { }) it('should be able to reset and run a new checkpoint experiment from an existing one with dvc.views.experiments.resetAndRunCheckpointExperiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { dvcRunner, experiments, experimentsModel } = buildExperiments({ disposer: disposable }) diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index ae09aee4df..e526c6f606 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -116,6 +116,7 @@ export const buildExperiments = ({ mockGetCommitMessages, mockSelectBranches, mockUpdateExperimentsData, + pipeline, resourceLocator } } diff --git a/extension/src/test/suite/experiments/workspace.test.ts b/extension/src/test/suite/experiments/workspace.test.ts index 8540ff4675..5b9c0b1f0c 100644 --- a/extension/src/test/suite/experiments/workspace.test.ts +++ b/extension/src/test/suite/experiments/workspace.test.ts @@ -16,8 +16,7 @@ import { bypassProgressCloseDelay, closeAllEditors, getInputBoxEvent, - getTimeSafeDisposer, - mockDuration + getTimeSafeDisposer } from '../util' import { dvcDemoPath } from '../../util' import { @@ -36,7 +35,6 @@ import { Title } from '../../../vscode/title' import { join } from '../../util/path' import { AvailableCommands } from '../../../commands/internal' import { EXPERIMENT_WORKSPACE_ID } from '../../../cli/dvc/contract' -import { DvcReader } from '../../../cli/dvc/reader' import { Setup } from '../../../setup' import { WorkspaceExperiments } from '../../../experiments/workspace' @@ -90,7 +88,7 @@ suite('Workspace Experiments Test Suite', () => { expect(mockQuickPickOne).to.be.calledOnce }).timeout(WEBVIEW_TEST_TIMEOUT) - it.only('should not prompt to pick a project if there is only one project', async () => { + it('should not prompt to pick a project if there is only one project', async () => { const mockQuickPickOne = stub(QuickPick, 'quickPickOne').resolves( dvcDemoPath ) @@ -111,7 +109,6 @@ suite('Workspace Experiments Test Suite', () => { DvcRunner.prototype, 'runExperiment' ).resolves(undefined) - stub(DvcReader.prototype, 'stageList').resolves('train') const { workspaceExperiments, experiments } = buildMultiRepoExperiments(disposable) @@ -160,8 +157,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndQueue', () => { it('should be able to queue an experiment using an existing one as a base', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { dvcExecutor, experiments } = buildExperiments({ disposer: disposable }) @@ -210,8 +205,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndResume', () => { it('should be able to resume a checkpoint experiment using an existing one as a base', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { experiments } = buildExperiments({ disposer: disposable }) const mockExperimentRun = stub( @@ -262,8 +255,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndRun', () => { it('should be able to run an experiment using an existing one as a base', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { experiments } = buildExperiments({ disposer: disposable }) const mockExperimentRun = stub( @@ -314,8 +305,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.queueExperiment', () => { it('should be able to queue an experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const mockExperimentRunQueue = stub( DvcExecutor.prototype, 'expRunQueue' @@ -330,8 +319,6 @@ suite('Workspace Experiments Test Suite', () => { }) it('should send a telemetry event containing a duration when an experiment is queued', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - stub(DvcExecutor.prototype, 'expRunQueue').resolves('true') const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent') @@ -352,11 +339,6 @@ suite('Workspace Experiments Test Suite', () => { }) it('should send a telemetry event containing an error message when an experiment fails to queue', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - - const duration = 77777 - mockDuration(duration) - const mockErrorMessage = 'ERROR: unexpected error - [Errno 2] No such file or directory' @@ -376,8 +358,7 @@ suite('Workspace Experiments Test Suite', () => { expect(mockSendTelemetryEvent).to.be.calledWith( `errors.${RegisteredCliCommands.QUEUE_EXPERIMENT}`, - { error: mockErrorMessage }, - { duration } + { error: mockErrorMessage } ) expect(mockGenericError, 'the generic error should be shown').to.be .calledOnce @@ -386,7 +367,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.runExperiment', () => { it('should be able to run an experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') const mockRunExperiment = stub( DvcRunner.prototype, 'runExperiment' @@ -407,8 +387,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.resumeCheckpointExperiment', () => { it('should be able to run an experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const mockRunExperiment = stub( DvcRunner.prototype, 'runExperiment' @@ -425,7 +403,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.resetAndRunCheckpointExperiment', () => { it('should be able to reset existing checkpoints and restart the experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') const mockRunExperimentReset = stub( DvcRunner.prototype, 'runExperimentReset' @@ -504,8 +481,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.startExperimentsQueue', () => { it('should be able to start the experiments queue with the selected number of workers', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const mockQueueStart = stub(DvcExecutor.prototype, 'queueStart').resolves( undefined ) @@ -531,8 +506,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.stopExperimentsQueue', () => { it('should be able to stop the experiments queue', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const mockQueueStop = stub(DvcExecutor.prototype, 'queueStop').resolves( undefined ) @@ -649,8 +622,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.branchExperiment', () => { it('should be able to create a branch from an experiment', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { experiments } = buildExperiments({ disposer: disposable }) await experiments.isReady() @@ -898,8 +869,6 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.removeExperimentQueue', () => { it('should remove all queued experiments from the selected repository', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const { experiments } = buildExperiments({ disposer: disposable }) await experiments.isReady() From be234e022cf511ff5b760850debe94802481f8d8 Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Thu, 13 Jul 2023 12:43:29 +1000 Subject: [PATCH 04/12] move check or add pipeline into pipeline class (no-verify) --- extension/src/experiments/index.ts | 13 +- extension/src/experiments/webview/contract.ts | 1 - extension/src/experiments/webview/messages.ts | 7 +- extension/src/experiments/workspace.test.ts | 11 +- extension/src/experiments/workspace.ts | 175 +++--------------- extension/src/fileSystem/index.test.ts | 30 +-- extension/src/pipeline/data.ts | 6 +- extension/src/pipeline/index.ts | 57 ++++-- extension/src/pipeline/model/collect.test.ts | 25 +-- extension/src/pipeline/model/collect.ts | 22 +-- extension/src/pipeline/model/index.ts | 26 +-- .../test/fixtures/expShow/base/tableData.ts | 1 - .../src/test/suite/experiments/index.test.ts | 43 +---- .../test/suite/experiments/workspace.test.ts | 3 + extension/src/test/suite/pipeline/util.ts | 7 +- .../src/experiments/components/AddStage.tsx | 14 +- .../src/experiments/components/App.test.tsx | 20 -- webview/src/experiments/components/App.tsx | 4 - .../experiments/components/Experiments.tsx | 3 +- .../src/experiments/state/tableDataSlice.ts | 5 - webview/src/stories/Table.stories.tsx | 1 - 21 files changed, 146 insertions(+), 328 deletions(-) diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 161415ecf0..68723023f4 100644 --- a/extension/src/experiments/index.ts +++ b/extension/src/experiments/index.ts @@ -431,7 +431,8 @@ export class Experiments extends BaseRepository { } public async modifyWorkspaceParamsAndRun( - commandId: ModifiedExperimentAndRunCommandId + commandId: ModifiedExperimentAndRunCommandId, + cwd: string ) { const paramsToModify = await this.pickAndModifyParams() if (!paramsToModify) { @@ -440,13 +441,13 @@ export class Experiments extends BaseRepository { await this.internalCommands.executeCommand( commandId, - this.dvcRoot, + cwd, ...paramsToModify ) return this.notifyChanged() } - public async modifyWorkspaceParamsAndQueue() { + public async modifyWorkspaceParamsAndQueue(cwd: string) { const paramsToModify = await this.pickAndModifyParams() if (!paramsToModify) { return @@ -455,7 +456,7 @@ export class Experiments extends BaseRepository { await Toast.showOutput( this.internalCommands.executeCommand( AvailableCommands.EXP_QUEUE, - this.dvcRoot, + cwd, ...paramsToModify ) ) @@ -517,6 +518,10 @@ export class Experiments extends BaseRepository { return this.columns.getRelativeMetricsFiles() } + public getPipelineCwd() { + return this.pipeline.getCwd() + } + protected sendInitialWebviewData() { return this.webviewMessages.sendWebviewMessage() } diff --git a/extension/src/experiments/webview/contract.ts b/extension/src/experiments/webview/contract.ts index 23cb4fd54c..256f1acbfc 100644 --- a/extension/src/experiments/webview/contract.ts +++ b/extension/src/experiments/webview/contract.ts @@ -104,7 +104,6 @@ export type TableData = { hasConfig: boolean hasMoreCommits: Record hasRunningWorkspaceExperiment: boolean - hasValidDvcYaml: boolean isShowingMoreCommits: Record rows: Commit[] selectedForPlotsCount: number diff --git a/extension/src/experiments/webview/messages.ts b/extension/src/experiments/webview/messages.ts index 3985bcdca9..e1b2e0f887 100644 --- a/extension/src/experiments/webview/messages.ts +++ b/extension/src/experiments/webview/messages.ts @@ -251,11 +251,10 @@ export class WebviewMessages { this.experiments.getAvailableBranchesToShow().length > 0, hasCheckpoints: this.experiments.hasCheckpoints(), hasColumns: this.columns.hasNonDefaultColumns(), - hasConfig: !!this.pipeline.hasStage(), + hasConfig: !!(this.pipeline.hasStage() || this.pipeline.hasPipeline()), hasMoreCommits: this.experiments.getHasMoreCommits(), hasRunningWorkspaceExperiment: this.experiments.hasRunningWorkspaceExperiment(), - hasValidDvcYaml: !this.pipeline.hasInvalidRootDvcYaml(), isShowingMoreCommits: this.experiments.getIsShowingMoreCommits(), rows: this.experiments.getRowData(), selectedForPlotsCount: this.experiments.getSelectedRevisions().length, @@ -263,8 +262,8 @@ export class WebviewMessages { } } - private async addConfiguration() { - await this.pipeline.checkOrAddPipeline() + private addConfiguration() { + this.pipeline.checkOrAddPipeline() } private async setMaxTableHeadDepth() { diff --git a/extension/src/experiments/workspace.test.ts b/extension/src/experiments/workspace.test.ts index 6b14d2d780..c5610e3388 100644 --- a/extension/src/experiments/workspace.test.ts +++ b/extension/src/experiments/workspace.test.ts @@ -1,6 +1,6 @@ import { Disposable, Disposer } from '@hediet/std/disposable' import { Experiments } from '.' -import { scriptCommand, WorkspaceExperiments } from './workspace' +import { WorkspaceExperiments } from './workspace' import { quickPickManyValues, quickPickOne, @@ -23,6 +23,7 @@ import { } from '../fileSystem' import { Toast } from '../vscode/toast' import { pickFile } from '../vscode/resourcePicker' +import { ScriptCommand } from '../pipeline' const mockedShowWebview = jest.fn() const mockedDisposable = jest.mocked(Disposable) @@ -96,11 +97,13 @@ describe('Experiments', () => { { '/my/dvc/root': { getDvcRoot: () => mockedDvcRoot, + getPipelineCwd: () => mockedDvcRoot, pickCommitOrExperiment: mockedPickCommitOrExperiment, showWebview: mockedShowWebview } as unknown as Experiments, '/my/fun/dvc/root': { getDvcRoot: () => mockedOtherDvcRoot, + getPipelineCwd: () => mockedOtherDvcRoot, pickExperiment: jest.fn(), showWebview: jest.fn() } as unknown as Experiments @@ -427,7 +430,7 @@ describe('Experiments', () => { mockedDvcRoot, trainingScript, 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, false ) }) @@ -468,7 +471,7 @@ describe('Experiments', () => { mockedDvcRoot, 'path/to/training_script.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, false ) }) @@ -488,7 +491,7 @@ describe('Experiments', () => { mockedDvcRoot, 'path/to/training_script.ipynb', 'train', - scriptCommand.JUPYTER, + ScriptCommand.JUPYTER, false ) }) diff --git a/extension/src/experiments/workspace.ts b/extension/src/experiments/workspace.ts index 8f7d5bf511..6531c60c9e 100644 --- a/extension/src/experiments/workspace.ts +++ b/extension/src/experiments/workspace.ts @@ -16,39 +16,13 @@ import { import { ResourceLocator } from '../resourceLocator' import { Setup } from '../setup' import { Toast } from '../vscode/toast' -import { - getInput, - getPositiveIntegerInput, - getValidInput -} from '../vscode/inputBox' +import { getInput, getPositiveIntegerInput } from '../vscode/inputBox' import { BaseWorkspaceWebviews } from '../webview/workspace' import { Title } from '../vscode/title' import { ContextKey, setContextValue } from '../vscode/context' -import { - findOrCreateDvcYamlFile, - getFileExtension, - hasDvcYamlFile -} from '../fileSystem' -import { quickPickManyValues, quickPickOneOrInput } from '../vscode/quickPick' -import { pickFile } from '../vscode/resourcePicker' +import { quickPickManyValues } from '../vscode/quickPick' import { WorkspacePipeline } from '../pipeline/workspace' -export enum scriptCommand { - JUPYTER = 'jupyter nbconvert --to notebook --inplace --execute', - PYTHON = 'python' -} - -const getScriptCommand = (script: string) => { - switch (getFileExtension(script)) { - case '.py': - return scriptCommand.PYTHON - case '.ipynb': - return scriptCommand.JUPYTER - default: - return '' - } -} - export class WorkspaceExperiments extends BaseWorkspaceWebviews< Experiments, TableData @@ -179,43 +153,25 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< commandId: ModifiedExperimentAndRunCommandId, overrideRoot?: string ) { - const cwd = await this.getDvcRoot(overrideRoot) - if (!cwd) { - return - } - - const repository = this.getRepository(cwd) - if (!repository) { - return - } - - const shouldContinue = await this.checkOrAddPipeline(cwd) - if (!shouldContinue) { + const { cwd, repository } = await this.getRepositoryAndCwd(overrideRoot) + if (!(cwd && repository)) { return } - return await repository.modifyWorkspaceParamsAndRun(commandId) + return await repository.modifyWorkspaceParamsAndRun(commandId, cwd) } public async modifyWorkspaceParamsAndQueue(overrideRoot?: string) { - const cwd = await this.getDvcRoot(overrideRoot) - if (!cwd) { - return - } - const shouldContinue = await this.checkOrAddPipeline(cwd) // repository.getPipeLineCwd() -> !cwd return - if (!shouldContinue) { - return - } - const repository = this.getRepository(cwd) - if (!repository) { + const { cwd, repository } = await this.getRepositoryAndCwd(overrideRoot) + if (!(cwd && repository)) { return } - return await repository.modifyWorkspaceParamsAndQueue() + return await repository.modifyWorkspaceParamsAndQueue(cwd) } public async getCwdThenRun(commandId: CommandId) { - const cwd = await this.shouldRun() + const cwd = await this.getCwd() if (!cwd) { return 'Could not run task' @@ -238,7 +194,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< commandId: CommandId, quickPick: () => Thenable ) { - const cwd = await this.shouldRun() + const cwd = await this.getCwd() if (!cwd) { return } @@ -250,7 +206,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< } public async createExperimentBranch() { - const cwd = await this.shouldRun() + const cwd = await this.getCwd() if (!cwd) { return } @@ -287,7 +243,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< title: Title, options: { prompt: string; value: string } ) { - const cwd = await this.shouldRun() + const cwd = await this.getCwd() if (!cwd) { return } @@ -451,106 +407,19 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< return this.getRepository(dvcRoot)[method]() } - private async shouldRun() { - const cwd = await this.getFocusedOrOnlyOrPickProject() - if (!cwd) { - return - } - const shouldContinue = await this.checkOrAddPipeline(cwd) // repository.getPipeLineCwd() -> !cwd return - if (!shouldContinue) { - return - } - return cwd - } - - private async checkOrAddPipeline(cwd: string) { - const stages = await this.internalCommands.executeCommand( - AvailableCommands.STAGE_LIST, - cwd - ) - - if (hasDvcYamlFile(cwd) && stages === undefined) { - await Toast.showError( - 'Cannot perform task. Your dvc.yaml file is invalid.' - ) - return false - } - - if (!stages) { - return this.addPipeline(cwd) + private async getRepositoryAndCwd(overrideRoot?: string) { + const project = await this.getDvcRoot(overrideRoot) + if (!project) { + return { cwd: undefined, repository: undefined } } - return true + const repository = this.getRepository(project) + const cwd = await repository.getPipelineCwd() + return { cwd, repository } } - private async addPipeline(cwd: string) { - const stageName = await this.askForStageName() - if (!stageName) { - return false - } - - const { trainingScript, command, enteredManually } = - await this.askForTrainingScript() - if (!trainingScript) { - return false - } - void findOrCreateDvcYamlFile( - cwd, - trainingScript, - stageName, - command, - !enteredManually - ) - return true - } - - private async askForStageName() { - return await getValidInput( - Title.ENTER_STAGE_NAME, - (stageName?: string) => { - if (!stageName) { - return 'Stage name must not be empty' - } - if (!/^[a-z]/i.test(stageName)) { - return 'Stage name should start with a letter' - } - return /^\w+$/.test(stageName) - ? null - : 'Stage name should only include letters and numbers' - }, - { value: 'train' } - ) - } - - private async askForTrainingScript() { - const selectValue = 'select' - const pathOrSelect = await quickPickOneOrInput( - [{ label: 'Select from file explorer', value: selectValue }], - { - defaultValue: '', - placeholder: 'Path to script', - title: Title.ENTER_PATH_OR_CHOOSE_FILE - } - ) - - const trainingScript = - pathOrSelect === selectValue - ? await pickFile(Title.SELECT_TRAINING_SCRIPT) - : pathOrSelect - - if (!trainingScript) { - return { - command: undefined, - enteredManually: false, - trainingScript: undefined - } - } - - const command = - getScriptCommand(trainingScript) || - (await getInput(Title.ENTER_COMMAND_TO_RUN)) || - '' - const enteredManually = pathOrSelect !== selectValue - return { command, enteredManually, trainingScript } + private async getCwd(overrideRoot?: string) { + const { cwd } = await this.getRepositoryAndCwd(overrideRoot) + return cwd } private async pickExpThenRun( diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts index bbff1bb29e..2887fa95d6 100644 --- a/extension/src/fileSystem/index.test.ts +++ b/extension/src/fileSystem/index.test.ts @@ -19,7 +19,7 @@ import { } from '.' import { dvcDemoPath } from '../test/util' import { DOT_DVC } from '../cli/dvc/constants' -import { scriptCommand } from '../experiments/workspace' +import { ScriptCommand } from '../pipeline' jest.mock('../cli/dvc/reader') jest.mock('fs-extra', () => { @@ -219,7 +219,7 @@ describe('findOrCreateDvcYamlFile', () => { cwd, '/my/training/script.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) @@ -233,7 +233,7 @@ describe('findOrCreateDvcYamlFile', () => { cwd, '/script.py', uniqueStageName, - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) @@ -249,7 +249,7 @@ describe('findOrCreateDvcYamlFile', () => { cwd, '/my/training/script.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) @@ -265,7 +265,7 @@ describe('findOrCreateDvcYamlFile', () => { cwd, '/my/training/script.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) @@ -282,7 +282,7 @@ describe('findOrCreateDvcYamlFile', () => { '/dir/my_project/', '/dir/my_project/src/training/train.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) @@ -295,7 +295,7 @@ describe('findOrCreateDvcYamlFile', () => { '/dir/my_project/', '/dir/my_other_project/train.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) @@ -312,7 +312,7 @@ describe('findOrCreateDvcYamlFile', () => { '/dir/my_project/', join('dir', 'my_project', 'src', 'training', 'train.py'), 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, false ) @@ -329,17 +329,17 @@ describe('findOrCreateDvcYamlFile', () => { '/', '/train.ipynb', 'train', - scriptCommand.JUPYTER, + ScriptCommand.JUPYTER, true ) expect(mockedAppendFileSync).toHaveBeenCalledWith( expect.anything(), - expect.stringContaining(scriptCommand.JUPYTER) + expect.stringContaining(ScriptCommand.JUPYTER) ) expect(mockedAppendFileSync).not.toHaveBeenCalledWith( expect.anything(), - expect.stringContaining(scriptCommand.PYTHON) + expect.stringContaining(ScriptCommand.PYTHON) ) }) @@ -348,17 +348,17 @@ describe('findOrCreateDvcYamlFile', () => { '/', '/train.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) expect(mockedAppendFileSync).not.toHaveBeenCalledWith( expect.anything(), - expect.stringContaining(scriptCommand.JUPYTER) + expect.stringContaining(ScriptCommand.JUPYTER) ) expect(mockedAppendFileSync).toHaveBeenCalledWith( expect.anything(), - expect.stringContaining(scriptCommand.PYTHON) + expect.stringContaining(ScriptCommand.PYTHON) ) }) @@ -379,7 +379,7 @@ describe('findOrCreateDvcYamlFile', () => { '/', '/train.py', 'train', - scriptCommand.PYTHON, + ScriptCommand.PYTHON, true ) diff --git a/extension/src/pipeline/data.ts b/extension/src/pipeline/data.ts index b7c33cc0a2..90a2bf1b8e 100644 --- a/extension/src/pipeline/data.ts +++ b/extension/src/pipeline/data.ts @@ -25,7 +25,7 @@ export class PipelineData extends BaseData<{ public async update(): Promise { const [dag, fileList] = await Promise.all([ this.internalCommands.executeCommand(AvailableCommands.DAG, this.dvcRoot), - findFiles('**/dvc.yaml') + this.findDvcYamls() ]) const dvcYamlsDirs = new Set() @@ -45,4 +45,8 @@ export class PipelineData extends BaseData<{ return this.notifyChanged({ dag, stages }) } + + private findDvcYamls() { + return findFiles('**/dvc.yaml') + } } diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index aaec8e09f1..2eec95fc88 100644 --- a/extension/src/pipeline/index.ts +++ b/extension/src/pipeline/index.ts @@ -10,10 +10,10 @@ import { findOrCreateDvcYamlFile, getFileExtension } from '../fileSystem' import { Toast } from '../vscode/toast' import { getInput, getValidInput } from '../vscode/inputBox' import { Title } from '../vscode/title' -import { quickPickOneOrInput } from '../vscode/quickPick' +import { quickPickOne, quickPickOneOrInput } from '../vscode/quickPick' import { pickFile } from '../vscode/resourcePicker' -enum scriptCommand { +export enum ScriptCommand { JUPYTER = 'jupyter nbconvert --to notebook --inplace --execute', PYTHON = 'python' } @@ -21,9 +21,9 @@ enum scriptCommand { const getScriptCommand = (script: string) => { switch (getFileExtension(script)) { case '.py': - return scriptCommand.PYTHON + return ScriptCommand.PYTHON case '.ipynb': - return scriptCommand.JUPYTER + return ScriptCommand.JUPYTER default: return '' } @@ -38,10 +38,16 @@ export class Pipeline extends DeferredDisposable { private readonly data: PipelineData private readonly model: PipelineModel - constructor(dvcRoot: string, internalCommands: InternalCommands) { + constructor( + dvcRoot: string, + internalCommands: InternalCommands, + data?: PipelineData + ) { super() this.dvcRoot = dvcRoot - this.data = this.dispose.track(new PipelineData(dvcRoot, internalCommands)) + this.data = this.dispose.track( + data || new PipelineData(dvcRoot, internalCommands) + ) this.model = this.dispose.track(new PipelineModel(dvcRoot)) this.updated = this.dispose.track(new EventEmitter()) this.onDidUpdate = this.updated.event @@ -53,18 +59,35 @@ export class Pipeline extends DeferredDisposable { return this.model.hasStage() } - public hasInvalidRootDvcYaml() { - return this.model.hasInvalidRootDvcYaml() + public hasPipeline() { + return this.model.hasPipeline() } - public async checkOrAddPipeline() { - const hasStage = this.model.hasStage() - if (this.hasInvalidRootDvcYaml()) { - await Toast.showError( - 'Cannot perform task. Your dvc.yaml file is invalid.' - ) - return false + public getCwd() { + const hasPipeline = this.checkOrAddPipeline() // need to refactor if we are going to run the original command + if (!hasPipeline) { + return + } + + 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' + ) + } + + public checkOrAddPipeline() { + const hasStage = this.model.hasStage() if (!hasStage) { return this.addPipeline() @@ -81,11 +104,11 @@ export class Pipeline extends DeferredDisposable { this.data.onDidUpdate(({ dag, stages }) => { this.writeDag(dag) const hasStage = this.model.hasStage() - const hasInvalidRootDvcYaml = this.model.hasInvalidRootDvcYaml() + const hasPipeline = this.model.hasPipeline() this.model.transformAndSet(stages) if ( hasStage !== this.model.hasStage() || - hasInvalidRootDvcYaml !== this.model.hasInvalidRootDvcYaml() + hasPipeline !== this.model.hasPipeline() ) { this.updated.fire() } diff --git a/extension/src/pipeline/model/collect.test.ts b/extension/src/pipeline/model/collect.test.ts index 6bb096a34b..3ebfbbba3d 100644 --- a/extension/src/pipeline/model/collect.test.ts +++ b/extension/src/pipeline/model/collect.test.ts @@ -2,37 +2,30 @@ import { join } from 'path' import { collectStages } from './collect' import { dvcDemoPath } from '../../test/util' -// need to combine collect stages and collect pipelines as a pipeline is only valid if it has stages - describe('collectStages', () => { it('should handle a simple list of stages', () => { - const { validPipelines, validStages } = collectStages({ + const { pipelines, stages } = collectStages({ [dvcDemoPath]: 'data.dvc Outputs data\ntrain Outputs model.pt, training/plots, hist.csv; Reports training/metrics.json' }) - expect(validStages).toStrictEqual(['train']) - expect(validPipelines).toStrictEqual(new Set([dvcDemoPath])) + expect(stages).toStrictEqual(['train']) + expect(pipelines).toStrictEqual(new Set([dvcDemoPath])) }) it('should handle a monorepos list of stages', () => { - const { validStages, validPipelines, invalidPipelines } = collectStages({ + const { stages, pipelines } = collectStages({ [join(dvcDemoPath, 'nested1')]: `data/data.xml.dvc Outputs data.xml prepare Outputs data/prepared featurize Outputs data/features train Outputs model.pkl evaluate Outputs eval/importance.png, eval/live/plots, eval/prc; Reports eval/live/metri…`, - [join('dvcDemoPath', 'nested2')]: 'data/data.xml.dvc Outputs data.xml' + [join('dvcDemoPath', 'nested2')]: 'data/data.xml.dvc Outputs data.xml', + [join('dvcDemoPath', 'nested3')]: undefined }) - expect(validStages).toStrictEqual([ - 'prepare', - 'featurize', - 'train', - 'evaluate' - ]) + expect(stages).toStrictEqual(['prepare', 'featurize', 'train', 'evaluate']) - expect(validPipelines).toStrictEqual( - new Set([join(dvcDemoPath, 'nested1')]) + expect(pipelines).toStrictEqual( + new Set([join(dvcDemoPath, 'nested1'), join(dvcDemoPath, 'nested3')]) ) - expect(invalidPipelines).toStrictEqual(new Set()) }) }) diff --git a/extension/src/pipeline/model/collect.ts b/extension/src/pipeline/model/collect.ts index e561053724..52e3facae6 100644 --- a/extension/src/pipeline/model/collect.ts +++ b/extension/src/pipeline/model/collect.ts @@ -1,20 +1,18 @@ import { trimAndSplit } from '../../util/stdout' -export const collectStages = (pipelines: { +export const collectStages = (data: { [pipeline: string]: string | undefined }): { - invalidPipelines: Set - validPipelines: Set - validStages: string[] + pipelines: Set + stages: string[] // eslint-disable-next-line sonarjs/cognitive-complexity } => { - const validStages: string[] = [] - const validPipelines = new Set() - const invalidPipelines = new Set() + const stages: string[] = [] + const pipelines = new Set() - for (const [pipeline, stageList] of Object.entries(pipelines)) { + for (const [pipeline, stageList] of Object.entries(data)) { if (stageList === undefined) { - invalidPipelines.add(pipeline) + pipelines.add(pipeline) } if (!stageList) { @@ -25,9 +23,9 @@ export const collectStages = (pipelines: { if (!stage || stage.endsWith('.dvc')) { continue } - validStages.push(stage) - validPipelines.add(pipeline) + stages.push(stage) + pipelines.add(pipeline) } } - return { invalidPipelines, validPipelines, validStages } + return { pipelines, stages } } diff --git a/extension/src/pipeline/model/index.ts b/extension/src/pipeline/model/index.ts index cbd567832c..449091c864 100644 --- a/extension/src/pipeline/model/index.ts +++ b/extension/src/pipeline/model/index.ts @@ -5,28 +5,24 @@ import { Disposable } from '../../class/dispose' export class PipelineModel extends Disposable { private readonly dvcRoot: string private stages: string[] | undefined = [] - private validPipelines: Set | undefined - private invalidPipelines: Set | undefined + private pipelines: Set | undefined constructor(dvcRoot: string) { super() this.dvcRoot = dvcRoot } - public transformAndSet(stages: { [pipeline: string]: string | undefined }) { - if (isEqual(stages, {})) { + public transformAndSet(data: { [pipeline: string]: string | undefined }) { + if (isEqual(data, {})) { this.stages = undefined - this.validPipelines = undefined - this.invalidPipelines = undefined + this.pipelines = undefined return } - const { validPipelines, invalidPipelines, validStages } = - collectStages(stages) + const { pipelines, stages } = collectStages(data) - this.validPipelines = validPipelines - this.invalidPipelines = invalidPipelines - this.stages = validStages + this.pipelines = pipelines + this.stages = stages } public hasStage() { @@ -38,14 +34,10 @@ export class PipelineModel extends Disposable { } public getPipelines() { - return this.validPipelines + return this.pipelines } public hasPipeline() { - return !!(this.validPipelines && this.validPipelines.size > 0) - } - - public hasInvalidRootDvcYaml() { - return !!this.invalidPipelines?.has(this.dvcRoot) + return !!(this.pipelines && this.pipelines.size > 0) } } diff --git a/extension/src/test/fixtures/expShow/base/tableData.ts b/extension/src/test/fixtures/expShow/base/tableData.ts index d01c0366e9..2c7ebbba1b 100644 --- a/extension/src/test/fixtures/expShow/base/tableData.ts +++ b/extension/src/test/fixtures/expShow/base/tableData.ts @@ -15,7 +15,6 @@ const tableDataFixture: TableData = { hasConfig: true, hasMoreCommits: { main: true }, hasRunningWorkspaceExperiment: true, - hasValidDvcYaml: true, isShowingMoreCommits: { main: true }, rows: rowsFixture, selectedForPlotsCount: 2, diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index f1edc6866a..0fd57414d6 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -139,7 +139,6 @@ suite('Experiments Test Suite', () => { hasColumns: true, hasConfig: true, hasRunningWorkspaceExperiment: true, - hasValidDvcYaml: true, rows: rowsFixture, sorts: [] } @@ -174,39 +173,7 @@ suite('Experiments Test Suite', () => { expect(windowSpy).not.to.have.been.called }).timeout(WEBVIEW_TEST_TIMEOUT) - it('should set hasValidDvcYaml to false if there is an error getting stages and there is a dvc.yaml file', async () => { - const { experiments, messageSpy } = buildExperiments({ - disposer: disposable, - stageList: null - }) - - await experiments.showWebview() - - expect(messageSpy).to.be.calledWithMatch({ - hasValidDvcYaml: false - }) - }).timeout(WEBVIEW_TEST_TIMEOUT) - - it('should set hasValidDvcYaml to true if there is no invalid root dvc.yaml file', async () => { - const { experiments, messageSpy, pipeline } = buildExperiments({ - disposer: disposable, - stageList: null - }) - - stub(pipeline, 'hasInvalidRootDvcYaml').returns(false) - - await experiments.showWebview() - - const expectedTableData = { - hasValidDvcYaml: true - } - - expect(messageSpy).to.be.calledWithMatch(expectedTableData) - }).timeout(WEBVIEW_TEST_TIMEOUT) - - it('should set hasValidDvcYaml to true if there are no errors getting stages', async () => { - stub(FileSystem, 'hasDvcYamlFile').returns(false) - + it('should set hasConfig to false if there are no stages', async () => { const { experiments, messageSpy } = buildExperiments({ disposer: disposable, stageList: '' @@ -215,20 +182,20 @@ suite('Experiments Test Suite', () => { await experiments.showWebview() expect(messageSpy).to.be.calledWithMatch({ - hasValidDvcYaml: true + hasConfig: false }) }).timeout(WEBVIEW_TEST_TIMEOUT) - it('should set hasConfig to false if there are no stages', async () => { + it('should set hasConfig to true if there is a broken dvc.yaml', async () => { const { experiments, messageSpy } = buildExperiments({ disposer: disposable, - stageList: '' + stageList: null }) await experiments.showWebview() expect(messageSpy).to.be.calledWithMatch({ - hasConfig: false + hasConfig: true }) }).timeout(WEBVIEW_TEST_TIMEOUT) diff --git a/extension/src/test/suite/experiments/workspace.test.ts b/extension/src/test/suite/experiments/workspace.test.ts index 5b9c0b1f0c..ebfbc9ebeb 100644 --- a/extension/src/test/suite/experiments/workspace.test.ts +++ b/extension/src/test/suite/experiments/workspace.test.ts @@ -160,6 +160,7 @@ suite('Workspace Experiments Test Suite', () => { const { dvcExecutor, experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() const mockExperimentRunQueue = stub(dvcExecutor, 'expRunQueue').resolves( 'true' @@ -206,6 +207,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndResume', () => { it('should be able to resume a checkpoint experiment using an existing one as a base', async () => { const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() const mockExperimentRun = stub( DvcRunner.prototype, @@ -256,6 +258,7 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.modifyWorkspaceParamsAndRun', () => { it('should be able to run an experiment using an existing one as a base', async () => { const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() const mockExperimentRun = stub( DvcRunner.prototype, diff --git a/extension/src/test/suite/pipeline/util.ts b/extension/src/test/suite/pipeline/util.ts index e395de7062..f8f14977a1 100644 --- a/extension/src/test/suite/pipeline/util.ts +++ b/extension/src/test/suite/pipeline/util.ts @@ -1,7 +1,9 @@ +import { join } from 'path' import { stub } from 'sinon' import { InternalCommands } from '../../../commands/internal' import { Disposer } from '../../../extension' import { Pipeline } from '../../../pipeline' +import { PipelineData } from '../../../pipeline/data' export const buildPipeline = ({ disposer, @@ -12,7 +14,10 @@ export const buildPipeline = ({ dvcRoot: string internalCommands: InternalCommands }): Pipeline => { - const pipeline = disposer.track(new Pipeline(dvcRoot, internalCommands)) + const data = new PipelineData(dvcRoot, internalCommands) + // eslint-disable-next-line @typescript-eslint/no-explicit-any + stub(data as any, 'findDvcYamls').resolves([join(dvcRoot, 'dvc.yaml')]) + const pipeline = disposer.track(new Pipeline(dvcRoot, internalCommands, data)) // eslint-disable-next-line @typescript-eslint/no-explicit-any stub(pipeline as any, 'writeDag').resolves() return pipeline diff --git a/webview/src/experiments/components/AddStage.tsx b/webview/src/experiments/components/AddStage.tsx index 9eb1652f74..a991ac3c03 100644 --- a/webview/src/experiments/components/AddStage.tsx +++ b/webview/src/experiments/components/AddStage.tsx @@ -4,11 +4,7 @@ import { IconButton } from '../../shared/components/button/IconButton' import { Add } from '../../shared/components/icons' import { addConfiguration } from '../util/messages' -interface AddStageProps { - hasValidDvcYaml: boolean -} - -export const AddStage: React.FC = ({ hasValidDvcYaml }) => ( +export const AddStage: React.FC = () => (

Define a{' '} @@ -19,15 +15,9 @@ export const AddStage: React.FC = ({ hasValidDvcYaml }) => (

hasValidDvcYaml && addConfiguration()} + onClick={() => addConfiguration()} text="Add Stage" - disabled={!hasValidDvcYaml} /> - {!hasValidDvcYaml && ( -

- A stage cannot be added to an invalid dvc.yaml file. -

- )}

Learn more about{' '} diff --git a/webview/src/experiments/components/App.test.tsx b/webview/src/experiments/components/App.test.tsx index 8c2d7d6d6c..bd738386e9 100644 --- a/webview/src/experiments/components/App.test.tsx +++ b/webview/src/experiments/components/App.test.tsx @@ -1759,26 +1759,6 @@ describe('App', () => { type: MessageFromWebviewType.ADD_CONFIGURATION }) }) - - it('should disable the button and add an error message if the dvc.yaml file contains invalid yaml', async () => { - renderTable() - setTableData({ - ...tableDataFixture, - hasConfig: false, - hasValidDvcYaml: false - }) - const addPipelineButton = await screen.findByText('Add Stage') - - fireEvent.click(addPipelineButton) - - expect(mockPostMessage).not.toHaveBeenCalledWith({ - type: MessageFromWebviewType.ADD_CONFIGURATION - }) - - expect( - screen.getByText('A stage cannot be added to an invalid dvc.yaml file.') - ).toBeInTheDocument() - }) }) describe('Show more commits', () => { diff --git a/webview/src/experiments/components/App.tsx b/webview/src/experiments/components/App.tsx index bf5a1446a5..7e1b665269 100644 --- a/webview/src/experiments/components/App.tsx +++ b/webview/src/experiments/components/App.tsx @@ -20,7 +20,6 @@ import { updateHasConfig, updateHasMoreCommits, updateHasRunningWorkspaceExperiment, - updateHasValidDvcYaml, updateIsShowingMoreCommits, updateRows, updateSelectedForPlotsCount, @@ -80,9 +79,6 @@ export const App: React.FC> = () => { ) ) continue - case 'hasValidDvcYaml': - dispatch(updateHasValidDvcYaml(data.data.hasValidDvcYaml)) - continue case 'isShowingMoreCommits': dispatch( updateIsShowingMoreCommits(data.data.isShowingMoreCommits) diff --git a/webview/src/experiments/components/Experiments.tsx b/webview/src/experiments/components/Experiments.tsx index d51fda5831..e8ffcee606 100644 --- a/webview/src/experiments/components/Experiments.tsx +++ b/webview/src/experiments/components/Experiments.tsx @@ -126,7 +126,6 @@ export const ExperimentsTable: React.FC = () => { columnWidths, hasColumns, hasConfig, - hasValidDvcYaml, rows: data } = useSelector((state: ExperimentsState) => state.tableData) @@ -186,7 +185,7 @@ export const ExperimentsTable: React.FC = () => { return ( - {!hasConfig && } + {!hasConfig && } ) } diff --git a/webview/src/experiments/state/tableDataSlice.ts b/webview/src/experiments/state/tableDataSlice.ts index d93c16a6a4..e38d31518d 100644 --- a/webview/src/experiments/state/tableDataSlice.ts +++ b/webview/src/experiments/state/tableDataSlice.ts @@ -28,7 +28,6 @@ export const tableDataInitialState: TableDataState = { hasData: false, hasMoreCommits: {}, hasRunningWorkspaceExperiment: false, - hasValidDvcYaml: true, isShowingMoreCommits: {}, rows: [], selectedForPlotsCount: 0, @@ -99,9 +98,6 @@ export const tableDataSlice = createSlice({ ) => { state.hasRunningWorkspaceExperiment = action.payload }, - updateHasValidDvcYaml: (state, action: PayloadAction) => { - state.hasValidDvcYaml = action.payload - }, updateIsShowingMoreCommits: ( state, action: PayloadAction> @@ -152,7 +148,6 @@ export const { updateHasConfig, updateHasMoreCommits, updateHasRunningWorkspaceExperiment, - updateHasValidDvcYaml, updateIsShowingMoreCommits, updateRows, updateSelectedForPlotsCount, diff --git a/webview/src/stories/Table.stories.tsx b/webview/src/stories/Table.stories.tsx index 273238ae7c..5e553b6563 100644 --- a/webview/src/stories/Table.stories.tsx +++ b/webview/src/stories/Table.stories.tsx @@ -48,7 +48,6 @@ const tableData: TableDataState = { hasData: true, hasMoreCommits: { main: true }, hasRunningWorkspaceExperiment: true, - hasValidDvcYaml: true, isShowingMoreCommits: { main: true }, rows: rowsFixture, selectedForPlotsCount: 2, From ad619be1423cda93769e153c4d9826b8438b884b Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Thu, 13 Jul 2023 16:34:05 +1000 Subject: [PATCH 05/12] move check or add pipeline tests into mocha --- extension/src/experiments/workspace.test.ts | 407 +----------------- extension/src/fileSystem/index.ts | 3 +- extension/src/pipeline/data.ts | 2 - extension/src/pipeline/index.ts | 35 +- extension/src/pipeline/model/collect.test.ts | 4 +- .../src/test/suite/experiments/index.test.ts | 7 +- extension/src/test/suite/experiments/util.ts | 30 +- .../src/test/suite/pipeline/index.test.ts | 292 +++++++++++++ extension/src/test/suite/pipeline/util.ts | 61 ++- extension/src/test/suite/plots/util.ts | 14 +- 10 files changed, 408 insertions(+), 447 deletions(-) create mode 100644 extension/src/test/suite/pipeline/index.test.ts diff --git a/extension/src/experiments/workspace.test.ts b/extension/src/experiments/workspace.test.ts index c5610e3388..66f13caf05 100644 --- a/extension/src/experiments/workspace.test.ts +++ b/extension/src/experiments/workspace.test.ts @@ -1,11 +1,7 @@ import { Disposable, Disposer } from '@hediet/std/disposable' import { Experiments } from '.' import { WorkspaceExperiments } from './workspace' -import { - quickPickManyValues, - quickPickOne, - quickPickOneOrInput -} from '../vscode/quickPick' +import { quickPickManyValues, quickPickOne } from '../vscode/quickPick' import { CommandId, AvailableCommands, @@ -16,14 +12,6 @@ import { buildMockMemento } from '../test/util' import { buildMockedEventEmitter } from '../test/util/jest' import { OutputChannel } from '../vscode/outputChannel' import { Title } from '../vscode/title' -import { - findOrCreateDvcYamlFile, - getFileExtension, - hasDvcYamlFile -} from '../fileSystem' -import { Toast } from '../vscode/toast' -import { pickFile } from '../vscode/resourcePicker' -import { ScriptCommand } from '../pipeline' const mockedShowWebview = jest.fn() const mockedDisposable = jest.mocked(Disposable) @@ -31,26 +19,19 @@ const mockedDvcRoot = '/my/dvc/root' const mockedOtherDvcRoot = '/my/fun/dvc/root' const mockedQuickPickOne = jest.mocked(quickPickOne) const mockedQuickPickManyValues = jest.mocked(quickPickManyValues) -const mockedQuickPickOneOrInput = jest.mocked(quickPickOneOrInput) const mockedGetValidInput = jest.mocked(getValidInput) const mockedPickCommitOrExperiment = jest.fn() const mockedGetInput = jest.mocked(getInput) const mockedRun = jest.fn() const mockedExpFunc = jest.fn() const mockedListStages = jest.fn() -const mockedFindOrCreateDvcYamlFile = jest.mocked(findOrCreateDvcYamlFile) -const mockedGetFileExtension = jest.mocked(getFileExtension) -const mockedHasDvcYamlFile = jest.mocked(hasDvcYamlFile) const mockedGetBranches = jest.fn() -const mockedPickFile = jest.mocked(pickFile) const mockedExpBranch = jest.fn() jest.mock('vscode') jest.mock('@hediet/std/disposable') jest.mock('../vscode/quickPick') jest.mock('../vscode/inputBox') -jest.mock('../fileSystem') -jest.mock('../vscode/resourcePicker') beforeEach(() => { jest.resetAllMocks() @@ -131,15 +112,6 @@ describe('Experiments', () => { expect(mockedQuickPickOne).toHaveBeenCalledTimes(1) expect(mockedExpFunc).not.toHaveBeenCalled() }) - - it('should check and ask for the creation of a pipeline stage before running the command', async () => { - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce('') - - await workspaceExperiments.getCwdThenReport(mockedCommandId) - - expect(mockedExpFunc).not.toHaveBeenCalled() - }) }) describe('getExpNameThenRun', () => { @@ -164,15 +136,6 @@ describe('Experiments', () => { expect(mockedQuickPickOne).toHaveBeenCalledTimes(1) expect(mockedExpFunc).not.toHaveBeenCalled() }) - - it('should check and ask for the creation of a pipeline stage before running the command', async () => { - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce('') - - await workspaceExperiments.getCwdAndExpNameThenRun(mockedCommandId) - - expect(mockedExpFunc).not.toHaveBeenCalled() - }) }) describe('getCwdAndQuickPickThenRun', () => { @@ -226,22 +189,6 @@ describe('Experiments', () => { expect(mockedQuickPick).toHaveBeenCalledTimes(1) expect(mockedExpFunc).not.toHaveBeenCalled() }) - - it('should check and ask for the creation of a pipeline stage before running the command', async () => { - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce('') - const mockedPickedOptions = ['a', 'b', 'c'] - const mockedQuickPick = jest - .fn() - .mockResolvedValueOnce(mockedPickedOptions) - - await workspaceExperiments.getCwdAndQuickPickThenRun( - mockedCommandId, - mockedQuickPick - ) - - expect(mockedExpFunc).not.toHaveBeenCalled() - }) }) describe('createExperimentBranch', () => { @@ -289,17 +236,6 @@ describe('Experiments', () => { expect(mockedGetInput).toHaveBeenCalledTimes(1) expect(mockedExpBranch).not.toHaveBeenCalled() }) - - it('should check and ask for the creation of a pipeline stage before doing anything else', async () => { - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce('') - mockedPickCommitOrExperiment.mockResolvedValueOnce('exp-123') - mockedGetInput.mockResolvedValueOnce('abc123') - - await workspaceExperiments.createExperimentBranch() - - expect(mockedExpBranch).not.toHaveBeenCalled() - }) }) describe('getCwdThenRun', () => { @@ -323,54 +259,22 @@ describe('Experiments', () => { expect(mockedExpFunc).not.toHaveBeenCalled() }) - it('should ensure that a dvc.yaml file exists', async () => { - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.py' - ) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledTimes(1) - }) - it('should check for pipelines when a command needs it and continue with the command if there is a pipeline', async () => { const executeCommandSpy = jest.spyOn( mockedInternalCommands, 'executeCommand' ) - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('train') mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) await workspaceExperiments.getCwdThenRun(mockedCommandId) - expect(executeCommandSpy).toHaveBeenCalledWith( - AvailableCommands.STAGE_LIST, - mockedDvcRoot - ) expect(executeCommandSpy).toHaveBeenCalledWith( mockedCommandId, mockedDvcRoot ) }) - it('should ask the user for the stage name if there are no pipelines', async () => { - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedGetValidInput).toHaveBeenCalledWith( - Title.ENTER_STAGE_NAME, - expect.anything(), - expect.anything() - ) - }) - it('should not ask the user for the stage name if there are pipelines', async () => { mockedListStages.mockResolvedValueOnce('train') mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) @@ -383,315 +287,6 @@ describe('Experiments', () => { expect.anything() ) }) - - it('should not run the command if no stage name was given', async () => { - const executeCommandSpy = jest.spyOn( - mockedInternalCommands, - 'executeCommand' - ) - - mockedGetValidInput.mockResolvedValueOnce('') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(executeCommandSpy).not.toHaveBeenCalledWith( - mockedCommandId, - mockedDvcRoot - ) - }) - - it('should let the user select a training script or enter its path if there are no pipelines found', async () => { - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.py' - ) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedQuickPickOneOrInput).toHaveBeenCalledTimes(1) - }) - - it('should add the train stage to the dvc.yaml file if the path to the training script was given', async () => { - const trainingScript = 'path/to/training_script.py' - - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce(trainingScript) - mockedGetFileExtension.mockReturnValueOnce('.py') - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith( - mockedDvcRoot, - trainingScript, - 'train', - ScriptCommand.PYTHON, - false - ) - }) - - it('should continue with the command if the path to the training script is entered', async () => { - const executeCommandSpy = jest.spyOn( - mockedInternalCommands, - 'executeCommand' - ) - - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.py' - ) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(executeCommandSpy).toHaveBeenCalledWith( - mockedCommandId, - mockedDvcRoot - ) - }) - - it('should add python as a command to the dvc.yaml file if the file has the .py extension', async () => { - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.py' - ) - mockedGetFileExtension.mockReturnValueOnce('.py') - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith( - mockedDvcRoot, - 'path/to/training_script.py', - 'train', - ScriptCommand.PYTHON, - false - ) - }) - - it('should add jupyter nbconvert as a command to the dvc.yaml file if the file has the .ipynb extension', async () => { - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.ipynb' - ) - mockedGetFileExtension.mockReturnValueOnce('.ipynb') - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith( - mockedDvcRoot, - 'path/to/training_script.ipynb', - 'train', - ScriptCommand.JUPYTER, - false - ) - }) - - it('should not ask to enter a custom command if the file is a python file or Jupyter notebook', async () => { - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.ipynb' - ) - mockedGetFileExtension.mockReturnValueOnce('.ipynb') - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedGetInput).not.toHaveBeenCalledWith( - Title.ENTER_COMMAND_TO_RUN - ) - }) - - it('should ask to enter a custom command if the file is not a python file or Jupyter notebook', async () => { - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.js' - ) - mockedGetFileExtension.mockReturnValueOnce('.js') - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedGetInput).toHaveBeenCalledWith(Title.ENTER_COMMAND_TO_RUN) - }) - - it('should add the custom command to the dvc.yaml file', async () => { - const customCommand = 'node' - - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.js' - ) - mockedGetFileExtension.mockReturnValueOnce('.js') - mockedGetInput.mockResolvedValueOnce(customCommand) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith( - mockedDvcRoot, - 'path/to/training_script.js', - 'train', - customCommand, - false - ) - }) - - it('should not convert the script path to relative if the path was entered manually', async () => { - const customCommand = 'node' - - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.js' - ) - mockedGetFileExtension.mockReturnValueOnce('.js') - mockedGetInput.mockResolvedValueOnce(customCommand) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith( - mockedDvcRoot, - 'path/to/training_script.js', - 'train', - customCommand, - false - ) - }) - - it('should convert the script path to relative if the path was not entered manually', async () => { - const customCommand = 'node' - - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce('select') - mockedPickFile.mockResolvedValueOnce('path/to/training_script.js') - mockedGetFileExtension.mockReturnValueOnce('.js') - mockedGetInput.mockResolvedValueOnce(customCommand) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith( - mockedDvcRoot, - 'path/to/training_script.js', - 'train', - customCommand, - true - ) - }) - - it('should not add a custom command to the dvc.yaml file if the command was not provided', async () => { - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce( - 'path/to/training_script.js' - ) - mockedGetFileExtension.mockReturnValueOnce('.js') - mockedGetInput.mockResolvedValueOnce(undefined) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith( - mockedDvcRoot, - 'path/to/training_script.js', - 'train', - '', - false - ) - }) - - it('should not run the command if the path to the training script was not given', async () => { - const executeCommandSpy = jest.spyOn( - mockedInternalCommands, - 'executeCommand' - ) - - mockedGetValidInput.mockResolvedValueOnce('train') - mockedListStages.mockResolvedValueOnce('') - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedQuickPickOneOrInput.mockResolvedValueOnce('') - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(executeCommandSpy).not.toHaveBeenCalledWith( - mockedCommandId, - mockedDvcRoot - ) - }) - - it('should not show a toast if there is no dvc.yaml file', async () => { - const showErrorSpy = jest.spyOn(Toast, 'showError') - - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce(undefined) - mockedHasDvcYamlFile.mockReturnValueOnce(false) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(showErrorSpy).not.toHaveBeenCalledWith( - 'Cannot perform task. Your dvc.yaml file is invalid.' - ) - }) - - it('should show a toast if the dvc.yaml file is invalid', async () => { - const showErrorSpy = jest.spyOn(Toast, 'showError') - - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce(undefined) - mockedHasDvcYamlFile.mockReturnValueOnce(true) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(showErrorSpy).toHaveBeenCalledWith( - 'Cannot perform task. Your dvc.yaml file is invalid.' - ) - }) - - it('should ask to create a stage if there is no dvc.yaml file', async () => { - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce(undefined) - mockedHasDvcYamlFile.mockReturnValueOnce(false) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedGetValidInput).toHaveBeenCalled() - }) - - it('should not ask to create a stage if the dvc.yaml file is invalid', async () => { - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce(undefined) - mockedHasDvcYamlFile.mockReturnValueOnce(true) - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(mockedGetValidInput).not.toHaveBeenCalled() - }) - - it('should not show a toast if the dvc.yaml file is valid', async () => { - const showErrorSpy = jest.spyOn(Toast, 'showError') - - mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot) - mockedListStages.mockResolvedValueOnce('train') - - await workspaceExperiments.getCwdThenRun(mockedCommandId) - - expect(showErrorSpy).not.toHaveBeenCalled() - }) }) describe('selectBranches', () => { diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts index 633260eee5..297fe02ee4 100644 --- a/extension/src/fileSystem/index.ts +++ b/extension/src/fileSystem/index.ts @@ -172,8 +172,7 @@ export const findOrCreateDvcYamlFile = ( ? relative(cwd, trainingScript) : format(parse(trainingScript)) - const pipeline = ` -# Type dvc-help in this file and hit enter to get more information on how the extension can help to setup pipelines + const pipeline = `# Type dvc-help in this file and hit enter to get more information on how the extension can help to setup pipelines stages: ${stageName}: cmd: ${command} ${scriptPath} diff --git a/extension/src/pipeline/data.ts b/extension/src/pipeline/data.ts index 90a2bf1b8e..2b976160ba 100644 --- a/extension/src/pipeline/data.ts +++ b/extension/src/pipeline/data.ts @@ -14,8 +14,6 @@ export class PipelineData extends BaseData<{ [{ name: 'update', process: () => this.update() }], ['dvc.yaml'] ) - - void this.managedUpdate() } public managedUpdate() { diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index 2eec95fc88..aeef2b051e 100644 --- a/extension/src/pipeline/index.ts +++ b/extension/src/pipeline/index.ts @@ -7,7 +7,6 @@ import { DeferredDisposable } from '../class/deferred' import { InternalCommands } from '../commands/internal' import { TEMP_DAG_FILE } from '../cli/dvc/constants' import { findOrCreateDvcYamlFile, getFileExtension } from '../fileSystem' -import { Toast } from '../vscode/toast' import { getInput, getValidInput } from '../vscode/inputBox' import { Title } from '../vscode/title' import { quickPickOne, quickPickOneOrInput } from '../vscode/quickPick' @@ -63,11 +62,8 @@ export class Pipeline extends DeferredDisposable { return this.model.hasPipeline() } - public getCwd() { - const hasPipeline = this.checkOrAddPipeline() // need to refactor if we are going to run the original command - if (!hasPipeline) { - return - } + public async getCwd() { + await this.checkOrAddPipeline() const pipelines = this.model.getPipelines() if (!pipelines?.size) { @@ -87,12 +83,10 @@ export class Pipeline extends DeferredDisposable { } public checkOrAddPipeline() { - const hasStage = this.model.hasStage() - - if (!hasStage) { - return this.addPipeline() + if (this.model.hasPipeline()) { + return } - return true + return this.addPipeline() } public forceRerender() { @@ -114,6 +108,7 @@ export class Pipeline extends DeferredDisposable { } }) ) + void this.data.managedUpdate() await this.data.isReady() return this.deferred.resolve() @@ -122,14 +117,25 @@ export class Pipeline extends DeferredDisposable { private async addPipeline() { const stageName = await this.askForStageName() if (!stageName) { - return false + return } const { trainingScript, command, enteredManually } = await this.askForTrainingScript() if (!trainingScript) { - return false + return } + + const dataUpdated = new Promise(resolve => { + const listener = this.dispose.track( + this.data.onDidUpdate(() => { + resolve(undefined) + this.dispose.untrack(listener) + listener.dispose() + }) + ) + }) + void findOrCreateDvcYamlFile( this.dvcRoot, trainingScript, @@ -137,7 +143,8 @@ export class Pipeline extends DeferredDisposable { command, !enteredManually ) - return true + void this.data.managedUpdate() + return dataUpdated } private async askForStageName() { diff --git a/extension/src/pipeline/model/collect.test.ts b/extension/src/pipeline/model/collect.test.ts index 3ebfbbba3d..bca2c50b88 100644 --- a/extension/src/pipeline/model/collect.test.ts +++ b/extension/src/pipeline/model/collect.test.ts @@ -19,8 +19,8 @@ describe('collectStages', () => { featurize Outputs data/features train Outputs model.pkl evaluate Outputs eval/importance.png, eval/live/plots, eval/prc; Reports eval/live/metri…`, - [join('dvcDemoPath', 'nested2')]: 'data/data.xml.dvc Outputs data.xml', - [join('dvcDemoPath', 'nested3')]: undefined + [join(dvcDemoPath, 'nested2')]: 'data/data.xml.dvc Outputs data.xml', + [join(dvcDemoPath, 'nested3')]: undefined }) expect(stages).toStrictEqual(['prepare', 'featurize', 'train', 'evaluate']) diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 0fd57414d6..9346b4960c 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -91,7 +91,7 @@ import { Toast } from '../../../vscode/toast' import { Response } from '../../../vscode/response' import { MAX_SELECTED_EXPERIMENTS } from '../../../experiments/model/status' import { Pipeline } from '../../../pipeline' -import { buildPipeline } from '../pipeline/util' +import { buildExperimentsPipeline } from '../pipeline/util' const { openFileInEditor } = FileSystem @@ -1510,14 +1510,15 @@ suite('Experiments Test Suite', () => { describe('Sorting', () => { it('should be able to sort', async () => { - const { internalCommands } = buildInternalCommands(disposable) + const { dvcReader, internalCommands } = buildInternalCommands(disposable) const messageSpy = spy(BaseWebview.prototype, 'show') const resourceLocator = disposable.track( new ResourceLocator(extensionUri) ) - const pipeline = buildPipeline({ + const pipeline = buildExperimentsPipeline({ + dvcReader, disposer: disposable, dvcRoot: dvcDemoPath, internalCommands diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index e526c6f606..84916ce916 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -19,7 +19,7 @@ import { ColumnsModel } from '../../../experiments/columns/model' import { DEFAULT_NUM_OF_COMMITS_TO_SHOW } from '../../../cli/dvc/constants' import { PersistenceKey } from '../../../persistence/constants' import { ExpShowOutput } from '../../../cli/dvc/contract' -import { buildPipeline } from '../pipeline/util' +import { buildExperimentsPipeline } from '../pipeline/util' export const DEFAULT_EXPERIMENTS_OUTPUT = { availableNbCommits: { main: 5 }, @@ -64,10 +64,13 @@ export const buildExperiments = ({ mockUpdateExperimentsData ) - stub(dvcReader, 'stageList').resolves(stageList ?? undefined) - stub(dvcReader, 'dag').resolves('') - - const pipeline = buildPipeline({ disposer, dvcRoot, internalCommands }) + const pipeline = buildExperimentsPipeline({ + disposer, + dvcReader, + dvcRoot, + internalCommands, + stageList + }) const mockCheckOrAddPipeline = stub(pipeline, 'checkOrAddPipeline') const mockSelectBranches = stub().resolves(['main', 'other']) const mockMemento = buildMockMemento({ @@ -123,6 +126,7 @@ export const buildExperiments = ({ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { const { + dvcReader, internalCommands, experiments: mockExperiments, gitReader, @@ -141,8 +145,9 @@ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { }) ) - const pipeline = buildPipeline({ + const pipeline = buildExperimentsPipeline({ disposer, + dvcReader, dvcRoot: dvcDemoPath, internalCommands }) @@ -159,15 +164,22 @@ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { } export const buildSingleRepoExperiments = (disposer: SafeWatcherDisposer) => { - const { config, internalCommands, gitReader, messageSpy, resourceLocator } = - buildDependencies(disposer) + const { + config, + dvcReader, + internalCommands, + gitReader, + messageSpy, + resourceLocator + } = buildDependencies(disposer) stub(gitReader, 'getGitRepositoryRoot').resolves(dvcDemoPath) const workspaceExperiments = disposer.track( new WorkspaceExperiments(internalCommands, buildMockMemento()) ) - const pipeline = buildPipeline({ + const pipeline = buildExperimentsPipeline({ disposer, + dvcReader, dvcRoot: dvcDemoPath, internalCommands }) diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts new file mode 100644 index 0000000000..093d60974f --- /dev/null +++ b/extension/src/test/suite/pipeline/index.test.ts @@ -0,0 +1,292 @@ +import { join } from 'path' +import { afterEach, beforeEach, describe, it, suite } from 'mocha' +import { SinonStub, restore, stub } from 'sinon' +import { expect } from 'chai' +import { QuickPickItem, Uri, window } from 'vscode' +import { buildPipeline } from './util' +import { + bypassProcessManagerDebounce, + closeAllEditors, + getMockNow +} from '../util' +import { Disposable } from '../../../extension' +import { dvcDemoPath } from '../../util' +import * as QuickPick from '../../../vscode/quickPick' +import { QuickPickOptionsWithTitle } from '../../../vscode/quickPick' +import * as FileSystem from '../../../fileSystem' +import { ScriptCommand } from '../../../pipeline' + +suite('Pipeline Test Suite', () => { + const disposable = Disposable.fn() + beforeEach(() => { + restore() + }) + + afterEach(() => { + return closeAllEditors() + }) + + describe('Pipeline', () => { + it('should ask to create a stage and not return a cwd if there is no pipeline', async () => { + const mockInputBox = stub(window, 'showInputBox').resolves(undefined) + const { pipeline } = buildPipeline({ + disposer: disposable, + dvcYamls: [] + }) + await pipeline.isReady() + const cwd = await pipeline.getCwd() + expect(mockInputBox, 'the user should be prompted to add a pipeline').to + .have.been.calledOnce + expect(cwd, 'no cwd is returned').to.be.undefined + }) + + it('should return a cwd if there is an invalid pipeline (let command fail)', async () => { + const { pipeline, pipelineModel } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + stageList: null + }) + await pipeline.isReady() + expect(pipelineModel.hasStage()).to.be.false + expect(pipelineModel.hasPipeline()).to.be.true + const cwd = await pipeline.getCwd() + expect(cwd).to.equal(dvcDemoPath) + }) + + it('should return a cwd if there is a single valid pipeline', async () => { + const { pipeline, pipelineModel } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + stageList: 'train' + }) + await pipeline.isReady() + expect(pipelineModel.hasStage()).to.be.true + expect(pipelineModel.hasPipeline()).to.be.true + const cwd = await pipeline.getCwd() + expect(cwd).to.equal(dvcDemoPath) + }) + + it('should return a the dvcRoot if there are multiple pipelines but one is the root', async () => { + const { pipeline, pipelineModel } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [dvcDemoPath, join(dvcDemoPath, 'subdir')], + stageList: 'train' + }) + await pipeline.isReady() + expect(pipelineModel.hasStage()).to.be.true + expect(pipelineModel.hasPipeline()).to.be.true + const cwd = await pipeline.getCwd() + expect(cwd).to.equal(dvcDemoPath) + }) + + it('should prompt the user to pick a pipeline if there are multiple pipelines and none are the root', async () => { + const pickedPipeline = join(dvcDemoPath, 'nested1', 'dvc.yaml') + const mockShowQuickPick = stub(window, 'showQuickPick') as SinonStub< + [items: readonly QuickPickItem[], options: QuickPickOptionsWithTitle], + Thenable + > + mockShowQuickPick.resolves(pickedPipeline) + const { pipeline, pipelineModel } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [pickedPipeline, join(dvcDemoPath, 'nested2', 'dvc.yaml')], + stageList: 'train' + }) + await pipeline.isReady() + expect(pipelineModel.hasStage()).to.be.true + expect(pipelineModel.hasPipeline()).to.be.true + const cwd = await pipeline.getCwd() + expect(cwd).to.equal(pickedPipeline) + expect(mockShowQuickPick).to.be.calledOnce + }) + + it('should create a new stage given all of the required information from the user', async () => { + const mockNow = getMockNow() + const { pipeline } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [] + }) + bypassProcessManagerDebounce(mockNow) + + const mockInputBox = stub(window, 'showInputBox').resolves('train') + const mockQuickPickOneOrInput = stub( + QuickPick, + 'quickPickOneOrInput' + ).resolves('train.py') + const mockFindOrCreateDvcYamlFile = stub( + FileSystem, + 'findOrCreateDvcYamlFile' + ).resolves(undefined) + + await pipeline.checkOrAddPipeline() + expect(mockInputBox).to.be.calledOnce + expect(mockQuickPickOneOrInput).to.be.calledOnce + expect(mockFindOrCreateDvcYamlFile).to.be.calledWithExactly( + dvcDemoPath, + 'train.py', + 'train', + 'python', + false + ) + }) + + it('should add jupyter nbconvert as a command to the dvc.yaml file if the file has the .ipynb extension', async () => { + const mockNow = getMockNow() + const { pipeline } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [] + }) + bypassProcessManagerDebounce(mockNow) + const scriptPath = join('path', 'to', 'training_script.ipynb') + const stageName = 'notebook_train' + + const mockInputBox = stub(window, 'showInputBox').resolves(stageName) + const mockQuickPickOneOrInput = stub( + QuickPick, + 'quickPickOneOrInput' + ).resolves(scriptPath) + const mockFindOrCreateDvcYamlFile = stub( + FileSystem, + 'findOrCreateDvcYamlFile' + ).resolves(undefined) + + await pipeline.checkOrAddPipeline() + expect(mockInputBox).to.be.calledOnce + expect(mockQuickPickOneOrInput).to.be.calledOnce + expect(mockFindOrCreateDvcYamlFile).to.be.calledWithExactly( + dvcDemoPath, + scriptPath, + stageName, + ScriptCommand.JUPYTER, + false + ) + }) + + it('should ask to enter a custom command if the file is not a python file or Jupyter notebook', async () => { + const mockNow = getMockNow() + const { pipeline } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [] + }) + bypassProcessManagerDebounce(mockNow) + const scriptCommand = 'go run' + const scriptPath = join('path', 'to', 'go-go-run.go') + const stageName = 'go-train' + + const mockInputBox = stub(window, 'showInputBox') + .onFirstCall() + .resolves(stageName) + .onSecondCall() + .resolves(scriptCommand) + const mockQuickPickOneOrInput = stub( + QuickPick, + 'quickPickOneOrInput' + ).resolves(scriptPath) + const mockFindOrCreateDvcYamlFile = stub( + FileSystem, + 'findOrCreateDvcYamlFile' + ).resolves(undefined) + + await pipeline.checkOrAddPipeline() + expect(mockInputBox).to.be.calledTwice + expect(mockQuickPickOneOrInput).to.be.calledOnce + expect(mockFindOrCreateDvcYamlFile).to.be.calledWithExactly( + dvcDemoPath, + scriptPath, + stageName, + scriptCommand, + false + ) + }) + + it('should ask to convert the script path to a relative path if the path was provided via the file picker', async () => { + const mockNow = getMockNow() + const { pipeline } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [] + }) + bypassProcessManagerDebounce(mockNow) + const scriptCommand = 'node' + const scriptPath = join(dvcDemoPath, 'path', 'to', 'training_script.js') + const stageName = 'whoTrainsInJavascript' + + const mockInputBox = stub(window, 'showInputBox') + .onFirstCall() + .resolves(stageName) + .onSecondCall() + .resolves(scriptCommand) + const mockQuickPickOneOrInput = stub( + QuickPick, + 'quickPickOneOrInput' + ).resolves('select') + const mockShowOpenDialog = stub(window, 'showOpenDialog').resolves([ + Uri.file(scriptPath) + ]) + + const mockFindOrCreateDvcYamlFile = stub( + FileSystem, + 'findOrCreateDvcYamlFile' + ).resolves(undefined) + + await pipeline.checkOrAddPipeline() + expect(mockInputBox).to.be.calledTwice + expect(mockQuickPickOneOrInput).to.be.calledOnce + expect(mockShowOpenDialog).to.be.calledOnce + expect(mockFindOrCreateDvcYamlFile).to.be.calledWithExactly( + dvcDemoPath, + scriptPath, + stageName, + scriptCommand, + true + ) + }) + + it('should not add a stage if a name was not given', async () => { + const { pipeline } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [] + }) + + const mockInputBox = stub(window, 'showInputBox').resolves(undefined) + const mockQuickPickOneOrInput = stub(QuickPick, 'quickPickOneOrInput') + const mockFindOrCreateDvcYamlFile = stub( + FileSystem, + 'findOrCreateDvcYamlFile' + ) + + await pipeline.checkOrAddPipeline() + expect(mockInputBox).to.be.calledOnce + expect(mockQuickPickOneOrInput).not.to.be.called + expect(mockFindOrCreateDvcYamlFile).not.to.be.called + }) + + it('should not add a stage if a training script was not given', async () => { + const { pipeline } = buildPipeline({ + disposer: disposable, + dvcRoot: dvcDemoPath, + dvcYamls: [] + }) + + const mockInputBox = stub(window, 'showInputBox').resolves('mega_train') + const mockQuickPickOneOrInput = stub( + QuickPick, + 'quickPickOneOrInput' + ).resolves(undefined) + const mockFindOrCreateDvcYamlFile = stub( + FileSystem, + 'findOrCreateDvcYamlFile' + ) + + await pipeline.checkOrAddPipeline() + expect(mockInputBox).to.be.calledOnce + expect(mockQuickPickOneOrInput).to.be.calledOnce + expect(mockFindOrCreateDvcYamlFile).not.to.be.called + }) + }) +}) diff --git a/extension/src/test/suite/pipeline/util.ts b/extension/src/test/suite/pipeline/util.ts index f8f14977a1..b2b14f5649 100644 --- a/extension/src/test/suite/pipeline/util.ts +++ b/extension/src/test/suite/pipeline/util.ts @@ -4,21 +4,72 @@ import { InternalCommands } from '../../../commands/internal' import { Disposer } from '../../../extension' import { Pipeline } from '../../../pipeline' import { PipelineData } from '../../../pipeline/data' +import { dvcDemoPath } from '../../util' +import { DvcReader } from '../../../cli/dvc/reader' +import { buildDependencies } from '../util' +import { PipelineModel } from '../../../pipeline/model' -export const buildPipeline = ({ +export const buildExperimentsPipeline = ({ + dag = '', disposer, - dvcRoot, - internalCommands + dvcRoot = dvcDemoPath, + dvcYamls, + dvcReader, + internalCommands, + stageList = 'train' }: { + dag?: string disposer: Disposer - dvcRoot: string + dvcRoot?: string + dvcYamls?: string[] + dvcReader: DvcReader internalCommands: InternalCommands + stageList?: string | null }): Pipeline => { + stub(dvcReader, 'stageList').resolves(stageList ?? undefined) + stub(dvcReader, 'dag').resolves(dag) + const data = new PipelineData(dvcRoot, internalCommands) // eslint-disable-next-line @typescript-eslint/no-explicit-any - stub(data as any, 'findDvcYamls').resolves([join(dvcRoot, 'dvc.yaml')]) + stub(data as any, 'findDvcYamls').resolves( + dvcYamls || [join(dvcRoot, 'dvc.yaml')] + ) const pipeline = disposer.track(new Pipeline(dvcRoot, internalCommands, data)) // eslint-disable-next-line @typescript-eslint/no-explicit-any stub(pipeline as any, 'writeDag').resolves() return pipeline } + +export const buildPipeline = ({ + dag = '', + disposer, + dvcRoot = dvcDemoPath, + dvcYamls, + stageList = 'train' +}: { + dag?: string + disposer: Disposer + dvcRoot?: string + dvcYamls?: string[] + stageList?: string | null +}) => { + const { dvcReader, internalCommands } = buildDependencies(disposer) + const pipeline = buildExperimentsPipeline({ + dag, + disposer, + dvcReader, + dvcRoot, + dvcYamls, + internalCommands, + stageList + }) + return { + dvcReader, + internalCommands, + pipeline, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + pipelineData: (pipeline as any).data as PipelineData, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + pipelineModel: (pipeline as any).model as PipelineModel + } +} diff --git a/extension/src/test/suite/plots/util.ts b/extension/src/test/suite/plots/util.ts index 8a21b64bbd..b95a89f7bf 100644 --- a/extension/src/test/suite/plots/util.ts +++ b/extension/src/test/suite/plots/util.ts @@ -26,7 +26,7 @@ import { } from '../../../cli/dvc/contract' import { ErrorsModel } from '../../../plots/errors/model' import { PersistenceKey } from '../../../persistence/constants' -import { buildPipeline } from '../pipeline/util' +import { buildExperimentsPipeline } from '../pipeline/util' export const buildPlots = async ({ availableNbCommits = { main: 5 }, @@ -43,16 +43,22 @@ export const buildPlots = async ({ gitLog?: string rowOrder?: { branch: string; sha: string }[] }) => { - const { internalCommands, mockPlotsDiff, messageSpy, resourceLocator } = - buildDependencies(disposer, expShow, plotsDiff) + const { + dvcReader, + internalCommands, + mockPlotsDiff, + messageSpy, + resourceLocator + } = buildDependencies(disposer, expShow, plotsDiff) const mockRemoveDir = stub(FileSystem, 'removeDir').returns(undefined) const mockGetModifiedTime = stub(FileSystem, 'getModifiedTime').returns( MOCK_IMAGE_MTIME ) - const pipeline = buildPipeline({ + const pipeline = buildExperimentsPipeline({ disposer, + dvcReader, dvcRoot: dvcDemoPath, internalCommands }) From 9928cb9ac01949c8d865923a9e8c149262d1ceda Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Thu, 13 Jul 2023 18:21:27 +1000 Subject: [PATCH 06/12] fix test pipeline builders --- extension/src/experiments/webview/messages.ts | 2 +- .../src/test/suite/experiments/index.test.ts | 4 +++- extension/src/test/suite/experiments/util.ts | 15 +++++++------- extension/src/test/suite/pipeline/util.ts | 20 +++++-------------- extension/src/test/suite/plots/util.ts | 4 +++- 5 files changed, 20 insertions(+), 25 deletions(-) diff --git a/extension/src/experiments/webview/messages.ts b/extension/src/experiments/webview/messages.ts index e1b2e0f887..1978bdc6d2 100644 --- a/extension/src/experiments/webview/messages.ts +++ b/extension/src/experiments/webview/messages.ts @@ -263,7 +263,7 @@ export class WebviewMessages { } private addConfiguration() { - this.pipeline.checkOrAddPipeline() + return this.pipeline.checkOrAddPipeline() } private async setMaxTableHeadDepth() { diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 9346b4960c..53bee3c193 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -1517,8 +1517,10 @@ suite('Experiments Test Suite', () => { const resourceLocator = disposable.track( new ResourceLocator(extensionUri) ) + stub(dvcReader, 'stageList').resolves('train') + stub(dvcReader, 'dag').resolves('') + const pipeline = buildExperimentsPipeline({ - dvcReader, disposer: disposable, dvcRoot: dvcDemoPath, internalCommands diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index 84916ce916..87b672e876 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -64,12 +64,13 @@ export const buildExperiments = ({ mockUpdateExperimentsData ) + stub(dvcReader, 'stageList').resolves(stageList ?? undefined) + stub(dvcReader, 'dag').resolves('') + const pipeline = buildExperimentsPipeline({ disposer, - dvcReader, dvcRoot, - internalCommands, - stageList + internalCommands }) const mockCheckOrAddPipeline = stub(pipeline, 'checkOrAddPipeline') const mockSelectBranches = stub().resolves(['main', 'other']) @@ -126,10 +127,9 @@ export const buildExperiments = ({ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { const { - dvcReader, - internalCommands, experiments: mockExperiments, gitReader, + internalCommands, messageSpy, resourceLocator } = buildExperiments({ @@ -147,7 +147,6 @@ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { const pipeline = buildExperimentsPipeline({ disposer, - dvcReader, dvcRoot: dvcDemoPath, internalCommands }) @@ -177,9 +176,11 @@ export const buildSingleRepoExperiments = (disposer: SafeWatcherDisposer) => { const workspaceExperiments = disposer.track( new WorkspaceExperiments(internalCommands, buildMockMemento()) ) + stub(dvcReader, 'stageList').resolves('train') + stub(dvcReader, 'dag').resolves('') + const pipeline = buildExperimentsPipeline({ disposer, - dvcReader, dvcRoot: dvcDemoPath, internalCommands }) diff --git a/extension/src/test/suite/pipeline/util.ts b/extension/src/test/suite/pipeline/util.ts index b2b14f5649..aec58da0ec 100644 --- a/extension/src/test/suite/pipeline/util.ts +++ b/extension/src/test/suite/pipeline/util.ts @@ -5,30 +5,20 @@ import { Disposer } from '../../../extension' import { Pipeline } from '../../../pipeline' import { PipelineData } from '../../../pipeline/data' import { dvcDemoPath } from '../../util' -import { DvcReader } from '../../../cli/dvc/reader' import { buildDependencies } from '../util' import { PipelineModel } from '../../../pipeline/model' export const buildExperimentsPipeline = ({ - dag = '', disposer, dvcRoot = dvcDemoPath, dvcYamls, - dvcReader, - internalCommands, - stageList = 'train' + internalCommands }: { - dag?: string disposer: Disposer dvcRoot?: string dvcYamls?: string[] - dvcReader: DvcReader internalCommands: InternalCommands - stageList?: string | null }): Pipeline => { - stub(dvcReader, 'stageList').resolves(stageList ?? undefined) - stub(dvcReader, 'dag').resolves(dag) - const data = new PipelineData(dvcRoot, internalCommands) // eslint-disable-next-line @typescript-eslint/no-explicit-any stub(data as any, 'findDvcYamls').resolves( @@ -54,14 +44,14 @@ export const buildPipeline = ({ stageList?: string | null }) => { const { dvcReader, internalCommands } = buildDependencies(disposer) + stub(dvcReader, 'stageList').resolves(stageList ?? undefined) + stub(dvcReader, 'dag').resolves(dag) + const pipeline = buildExperimentsPipeline({ - dag, disposer, - dvcReader, dvcRoot, dvcYamls, - internalCommands, - stageList + internalCommands }) return { dvcReader, diff --git a/extension/src/test/suite/plots/util.ts b/extension/src/test/suite/plots/util.ts index b95a89f7bf..50dbc56d52 100644 --- a/extension/src/test/suite/plots/util.ts +++ b/extension/src/test/suite/plots/util.ts @@ -56,9 +56,11 @@ export const buildPlots = async ({ MOCK_IMAGE_MTIME ) + stub(dvcReader, 'stageList').resolves('train') + stub(dvcReader, 'dag').resolves('') + const pipeline = buildExperimentsPipeline({ disposer, - dvcReader, dvcRoot: dvcDemoPath, internalCommands }) From d1727ce9e8fb695752abcc96328554bb8b125c0e Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Fri, 14 Jul 2023 09:04:14 +1000 Subject: [PATCH 07/12] stub experiments for all getters --- extension/src/test/suite/experiments/util.ts | 16 ++++---- .../test/suite/experiments/workspace.test.ts | 40 +++++++++++++++---- extension/src/test/suite/pipeline/util.ts | 2 +- 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index 87b672e876..d38362ff16 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -64,7 +64,7 @@ export const buildExperiments = ({ mockUpdateExperimentsData ) - stub(dvcReader, 'stageList').resolves(stageList ?? undefined) + stub(dvcReader, 'stageList').resolves(stageList ?? undefined) // move these two into buildDependencies stub(dvcReader, 'dag').resolves('') const pipeline = buildExperimentsPipeline({ @@ -256,19 +256,17 @@ export const buildExperimentsData = ( export const stubWorkspaceExperimentsGetters = ( dvcRoot: string, - experiments?: Experiments + experiments: Experiments ) => { const mockGetOnlyOrPickProject = stub( WorkspaceExperiments.prototype, 'getOnlyOrPickProject' ).resolves(dvcRoot) - let mockGetRepository - if (experiments) { - mockGetRepository = stub( - WorkspaceExperiments.prototype, - 'getRepository' - ).returns(experiments) - } + + const mockGetRepository = stub( + WorkspaceExperiments.prototype, + 'getRepository' + ).returns(experiments) return [mockGetOnlyOrPickProject, mockGetRepository] } diff --git a/extension/src/test/suite/experiments/workspace.test.ts b/extension/src/test/suite/experiments/workspace.test.ts index ebfbc9ebeb..2ffb6ec0f2 100644 --- a/extension/src/test/suite/experiments/workspace.test.ts +++ b/extension/src/test/suite/experiments/workspace.test.ts @@ -308,12 +308,15 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.queueExperiment', () => { it('should be able to queue an experiment', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + const mockExperimentRunQueue = stub( DvcExecutor.prototype, 'expRunQueue' ).resolves('true') - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) await commands.executeCommand(RegisteredCliCommands.QUEUE_EXPERIMENT) @@ -322,11 +325,14 @@ suite('Workspace Experiments Test Suite', () => { }) it('should send a telemetry event containing a duration when an experiment is queued', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + stub(DvcExecutor.prototype, 'expRunQueue').resolves('true') const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent') - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) const queueExperiment = commands.executeCommand( RegisteredCliCommands.QUEUE_EXPERIMENT @@ -342,6 +348,9 @@ suite('Workspace Experiments Test Suite', () => { }) it('should send a telemetry event containing an error message when an experiment fails to queue', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + const mockErrorMessage = 'ERROR: unexpected error - [Errno 2] No such file or directory' @@ -355,7 +364,7 @@ suite('Workspace Experiments Test Suite', () => { const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent') - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) await commands.executeCommand(RegisteredCliCommands.QUEUE_EXPERIMENT) @@ -370,6 +379,9 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.runExperiment', () => { it('should be able to run an experiment', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + const mockRunExperiment = stub( DvcRunner.prototype, 'runExperiment' @@ -379,7 +391,7 @@ suite('Workspace Experiments Test Suite', () => { experiments: true }) - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) await commands.executeCommand(RegisteredCliCommands.EXPERIMENT_RUN) @@ -390,12 +402,15 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.resumeCheckpointExperiment', () => { it('should be able to run an experiment', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + const mockRunExperiment = stub( DvcRunner.prototype, 'runExperiment' ).resolves(undefined) - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) await commands.executeCommand(RegisteredCliCommands.EXPERIMENT_RESUME) @@ -406,6 +421,9 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.resetAndRunCheckpointExperiment', () => { it('should be able to reset existing checkpoints and restart the experiment', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + const mockRunExperimentReset = stub( DvcRunner.prototype, 'runExperimentReset' @@ -415,7 +433,7 @@ suite('Workspace Experiments Test Suite', () => { experiments: true }) - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) await commands.executeCommand( RegisteredCliCommands.EXPERIMENT_RESET_AND_RUN @@ -484,6 +502,9 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.startExperimentsQueue', () => { it('should be able to start the experiments queue with the selected number of workers', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + const mockQueueStart = stub(DvcExecutor.prototype, 'queueStart').resolves( undefined ) @@ -494,7 +515,7 @@ suite('Workspace Experiments Test Suite', () => { dDosNumberOfJobs ) - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) await commands.executeCommand(RegisteredCliCommands.QUEUE_START) @@ -509,11 +530,14 @@ suite('Workspace Experiments Test Suite', () => { describe('dvc.stopExperimentsQueue', () => { it('should be able to stop the experiments queue', async () => { + const { experiments } = buildExperiments({ disposer: disposable }) + await experiments.isReady() + const mockQueueStop = stub(DvcExecutor.prototype, 'queueStop').resolves( undefined ) - stubWorkspaceExperimentsGetters(dvcDemoPath) + stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) await commands.executeCommand(RegisteredCliCommands.QUEUE_STOP) diff --git a/extension/src/test/suite/pipeline/util.ts b/extension/src/test/suite/pipeline/util.ts index aec58da0ec..787e89d901 100644 --- a/extension/src/test/suite/pipeline/util.ts +++ b/extension/src/test/suite/pipeline/util.ts @@ -44,7 +44,7 @@ export const buildPipeline = ({ stageList?: string | null }) => { const { dvcReader, internalCommands } = buildDependencies(disposer) - stub(dvcReader, 'stageList').resolves(stageList ?? undefined) + stub(dvcReader, 'stageList').resolves(stageList ?? undefined) // move into build dependencies stub(dvcReader, 'dag').resolves(dag) const pipeline = buildExperimentsPipeline({ From 4529c6ce2444ea4b0df5e5845f4ca4394dfb349b Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Fri, 14 Jul 2023 09:32:40 +1000 Subject: [PATCH 08/12] move reader stubs into build dependencies --- extension/src/test/suite/experiments/util.ts | 17 +++---------- extension/src/test/suite/extension.test.ts | 2 -- .../src/test/suite/pipeline/index.test.ts | 2 +- extension/src/test/suite/pipeline/util.ts | 8 ++++--- .../src/test/suite/plots/data/index.test.ts | 4 +++- extension/src/test/suite/plots/util.ts | 12 ++-------- .../test/suite/repository/model/tree.test.ts | 5 ++-- extension/src/test/suite/setup/util.ts | 2 +- extension/src/test/suite/util.ts | 24 +++++++++++++++---- 9 files changed, 38 insertions(+), 38 deletions(-) diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index d38362ff16..360af33b53 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -57,16 +57,13 @@ export const buildExperiments = ({ mockExpShow, mockGetCommitMessages, resourceLocator - } = buildDependencies(disposer, expShow) + } = buildDependencies({ disposer, expShow, stageList }) const mockUpdateExperimentsData = stub() const mockExperimentsData = buildMockExperimentsData( mockUpdateExperimentsData ) - stub(dvcReader, 'stageList').resolves(stageList ?? undefined) // move these two into buildDependencies - stub(dvcReader, 'dag').resolves('') - const pipeline = buildExperimentsPipeline({ disposer, dvcRoot, @@ -163,21 +160,13 @@ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { } export const buildSingleRepoExperiments = (disposer: SafeWatcherDisposer) => { - const { - config, - dvcReader, - internalCommands, - gitReader, - messageSpy, - resourceLocator - } = buildDependencies(disposer) + const { config, internalCommands, gitReader, messageSpy, resourceLocator } = + buildDependencies({ disposer }) stub(gitReader, 'getGitRepositoryRoot').resolves(dvcDemoPath) const workspaceExperiments = disposer.track( new WorkspaceExperiments(internalCommands, buildMockMemento()) ) - stub(dvcReader, 'stageList').resolves('train') - stub(dvcReader, 'dag').resolves('') const pipeline = buildExperimentsPipeline({ disposer, diff --git a/extension/src/test/suite/extension.test.ts b/extension/src/test/suite/extension.test.ts index 1277c3e9f9..eb41ef346a 100644 --- a/extension/src/test/suite/extension.test.ts +++ b/extension/src/test/suite/extension.test.ts @@ -219,8 +219,6 @@ suite('Extension Test Suite', () => { describe('dvc.stopAllRunningExperiments', () => { it('should send a telemetry event containing properties relating to the event', async () => { - stub(DvcReader.prototype, 'stageList').resolves('train') - const duration = 1234 const otherRoot = resolve('other', 'root') mockDuration(duration) diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts index 093d60974f..0809a86de5 100644 --- a/extension/src/test/suite/pipeline/index.test.ts +++ b/extension/src/test/suite/pipeline/index.test.ts @@ -66,7 +66,7 @@ suite('Pipeline Test Suite', () => { expect(cwd).to.equal(dvcDemoPath) }) - it('should return a the dvcRoot if there are multiple pipelines but one is the root', async () => { + it('should return the project root if there are multiple pipelines but one is the root', async () => { const { pipeline, pipelineModel } = buildPipeline({ disposer: disposable, dvcRoot: dvcDemoPath, diff --git a/extension/src/test/suite/pipeline/util.ts b/extension/src/test/suite/pipeline/util.ts index 787e89d901..dd5882eb46 100644 --- a/extension/src/test/suite/pipeline/util.ts +++ b/extension/src/test/suite/pipeline/util.ts @@ -43,9 +43,11 @@ export const buildPipeline = ({ dvcYamls?: string[] stageList?: string | null }) => { - const { dvcReader, internalCommands } = buildDependencies(disposer) - stub(dvcReader, 'stageList').resolves(stageList ?? undefined) // move into build dependencies - stub(dvcReader, 'dag').resolves(dag) + const { dvcReader, internalCommands } = buildDependencies({ + dag, + disposer, + stageList + }) const pipeline = buildExperimentsPipeline({ disposer, diff --git a/extension/src/test/suite/plots/data/index.test.ts b/extension/src/test/suite/plots/data/index.test.ts index e0368fa965..505931d33d 100644 --- a/extension/src/test/suite/plots/data/index.test.ts +++ b/extension/src/test/suite/plots/data/index.test.ts @@ -34,7 +34,9 @@ suite('Plots Data Test Suite', () => { }) const buildPlotsData = (selectedRevisions: string[] = []) => { - const { internalCommands, mockPlotsDiff } = buildDependencies(disposable) + const { internalCommands, mockPlotsDiff } = buildDependencies({ + disposer: disposable + }) const mockGetSelectedOrderedIds = stub().returns(selectedRevisions) diff --git a/extension/src/test/suite/plots/util.ts b/extension/src/test/suite/plots/util.ts index 50dbc56d52..f8ce99db33 100644 --- a/extension/src/test/suite/plots/util.ts +++ b/extension/src/test/suite/plots/util.ts @@ -43,22 +43,14 @@ export const buildPlots = async ({ gitLog?: string rowOrder?: { branch: string; sha: string }[] }) => { - const { - dvcReader, - internalCommands, - mockPlotsDiff, - messageSpy, - resourceLocator - } = buildDependencies(disposer, expShow, plotsDiff) + const { internalCommands, mockPlotsDiff, messageSpy, resourceLocator } = + buildDependencies({ disposer, expShow, plotsDiff }) const mockRemoveDir = stub(FileSystem, 'removeDir').returns(undefined) const mockGetModifiedTime = stub(FileSystem, 'getModifiedTime').returns( MOCK_IMAGE_MTIME ) - stub(dvcReader, 'stageList').resolves('train') - stub(dvcReader, 'dag').resolves('') - const pipeline = buildExperimentsPipeline({ disposer, dvcRoot: dvcDemoPath, diff --git a/extension/src/test/suite/repository/model/tree.test.ts b/extension/src/test/suite/repository/model/tree.test.ts index d2a3546132..168f3ff7c2 100644 --- a/extension/src/test/suite/repository/model/tree.test.ts +++ b/extension/src/test/suite/repository/model/tree.test.ts @@ -287,8 +287,9 @@ suite('Repositories Tree Test Suite', () => { }) it('should pull the correct target(s) when asked to dvc.pullTarget a non-tracked directory', async () => { - const { gitReader, internalCommands, mockDataStatus } = - buildDependencies(disposable) + const { gitReader, internalCommands, mockDataStatus } = buildDependencies( + { disposer: disposable } + ) mockDataStatus.resetBehavior() mockDataStatus.resolves({ diff --git a/extension/src/test/suite/setup/util.ts b/extension/src/test/suite/setup/util.ts index fe34bd1e8b..e4e59c23f5 100644 --- a/extension/src/test/suite/setup/util.ts +++ b/extension/src/test/suite/setup/util.ts @@ -47,7 +47,7 @@ export const buildSetup = ({ dvcReader, gitExecutor, gitReader - } = buildDependencies(disposer) + } = buildDependencies({ disposer }) if (gitVersion === undefined) { gitVersion = 'git version 2.41.0' diff --git a/extension/src/test/suite/util.ts b/extension/src/test/suite/util.ts index f621ac1741..abfe59feb0 100644 --- a/extension/src/test/suite/util.ts +++ b/extension/src/test/suite/util.ts @@ -37,6 +37,7 @@ import { DvcViewer } from '../../cli/dvc/viewer' import { Toast } from '../../vscode/toast' import { GitExecutor } from '../../cli/git/executor' import { DvcConfig } from '../../cli/dvc/config' +import { ExpShowOutput, PlotsOutput } from '../../cli/dvc/contract' export const mockDisposable = { dispose: stub() @@ -193,11 +194,19 @@ export const buildMockExperimentsData = (update = stub()) => const buildResourceLocator = (disposer: Disposer): ResourceLocator => disposer.track(new ResourceLocator(extensionUri)) -export const buildDependencies = ( - disposer: Disposer, +export const buildDependencies = ({ + dag = '', + disposer, expShow = expShowFixture, - plotsDiff = plotsDiffFixture -) => { + plotsDiff = plotsDiffFixture, + stageList = 'train' +}: { + dag?: string | undefined + disposer: Disposer + expShow?: ExpShowOutput + stageList?: string | null + plotsDiff?: PlotsOutput +}) => { const { config, dvcConfig, @@ -229,6 +238,11 @@ export const buildDependencies = ( '' ) + const mockStageList = stub(dvcReader, 'stageList').resolves( + stageList ?? undefined + ) + const mockDag = stub(dvcReader, 'dag').resolves(dag) + const resourceLocator = buildResourceLocator(disposer) const messageSpy = spy(BaseWebview.prototype, 'show') @@ -246,10 +260,12 @@ export const buildDependencies = ( messageSpy, mockCheckSignalFile, mockCreateFileSystemWatcher, + mockDag, mockDataStatus, mockExpShow, mockGetCommitMessages, mockPlotsDiff, + mockStageList, resourceLocator } } From b9736cf6e4d43db92c3580900046b88805225eb3 Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Fri, 14 Jul 2023 09:38:10 +1000 Subject: [PATCH 09/12] add stubs for on did update --- extension/src/test/suite/experiments/index.test.ts | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 53bee3c193..ccff8ad1ec 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -1707,8 +1707,9 @@ suite('Experiments Test Suite', () => { internalCommands, { hasStage: () => true, - isReady: () => Promise.resolve() - } as Pipeline, + isReady: () => Promise.resolve(), + onDidUpdate: stub() + } as unknown as Pipeline, {} as ResourceLocator, mockMemento, () => Promise.resolve([]), @@ -1871,8 +1872,9 @@ suite('Experiments Test Suite', () => { internalCommands, { hasStage: () => true, - isReady: () => Promise.resolve() - } as Pipeline, + isReady: () => Promise.resolve(), + onDidUpdate: stub() + } as unknown as Pipeline, {} as ResourceLocator, mockMemento, () => Promise.resolve([]), From 151dad74e1eea55bca8aa78c9a16f3fca37bb743 Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Fri, 14 Jul 2023 10:14:37 +1000 Subject: [PATCH 10/12] refactor pipeline collection --- extension/src/experiments/webview/messages.ts | 2 +- extension/src/pipeline/index.ts | 12 +---- extension/src/pipeline/model/collect.test.ts | 11 ++--- extension/src/pipeline/model/collect.ts | 47 ++++++++++--------- extension/src/pipeline/model/index.ts | 21 +-------- extension/src/test/suite/experiments/util.ts | 2 +- .../src/test/suite/pipeline/index.test.ts | 20 ++++---- extension/src/test/suite/pipeline/util.ts | 5 +- 8 files changed, 46 insertions(+), 74 deletions(-) diff --git a/extension/src/experiments/webview/messages.ts b/extension/src/experiments/webview/messages.ts index 1978bdc6d2..097f54dd8d 100644 --- a/extension/src/experiments/webview/messages.ts +++ b/extension/src/experiments/webview/messages.ts @@ -251,7 +251,7 @@ export class WebviewMessages { this.experiments.getAvailableBranchesToShow().length > 0, hasCheckpoints: this.experiments.hasCheckpoints(), hasColumns: this.columns.hasNonDefaultColumns(), - hasConfig: !!(this.pipeline.hasStage() || this.pipeline.hasPipeline()), + hasConfig: this.pipeline.hasPipeline(), hasMoreCommits: this.experiments.getHasMoreCommits(), hasRunningWorkspaceExperiment: this.experiments.hasRunningWorkspaceExperiment(), diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts index aeef2b051e..7b6486d43b 100644 --- a/extension/src/pipeline/index.ts +++ b/extension/src/pipeline/index.ts @@ -47,17 +47,13 @@ export class Pipeline extends DeferredDisposable { this.data = this.dispose.track( data || new PipelineData(dvcRoot, internalCommands) ) - this.model = this.dispose.track(new PipelineModel(dvcRoot)) + this.model = this.dispose.track(new PipelineModel()) this.updated = this.dispose.track(new EventEmitter()) this.onDidUpdate = this.updated.event void this.initialize() } - public hasStage() { - return this.model.hasStage() - } - public hasPipeline() { return this.model.hasPipeline() } @@ -97,13 +93,9 @@ export class Pipeline extends DeferredDisposable { this.dispose.track( this.data.onDidUpdate(({ dag, stages }) => { this.writeDag(dag) - const hasStage = this.model.hasStage() const hasPipeline = this.model.hasPipeline() this.model.transformAndSet(stages) - if ( - hasStage !== this.model.hasStage() || - hasPipeline !== this.model.hasPipeline() - ) { + if (hasPipeline !== this.model.hasPipeline()) { this.updated.fire() } }) diff --git a/extension/src/pipeline/model/collect.test.ts b/extension/src/pipeline/model/collect.test.ts index bca2c50b88..d6520cb1c7 100644 --- a/extension/src/pipeline/model/collect.test.ts +++ b/extension/src/pipeline/model/collect.test.ts @@ -1,19 +1,19 @@ import { join } from 'path' -import { collectStages } from './collect' +import { collectPipelines } from './collect' import { dvcDemoPath } from '../../test/util' -describe('collectStages', () => { +describe('collectPipelines', () => { it('should handle a simple list of stages', () => { - const { pipelines, stages } = collectStages({ + const pipelines = collectPipelines({ [dvcDemoPath]: 'data.dvc Outputs data\ntrain Outputs model.pt, training/plots, hist.csv; Reports training/metrics.json' }) - expect(stages).toStrictEqual(['train']) + expect(pipelines).toStrictEqual(new Set([dvcDemoPath])) }) it('should handle a monorepos list of stages', () => { - const { stages, pipelines } = collectStages({ + const pipelines = collectPipelines({ [join(dvcDemoPath, 'nested1')]: `data/data.xml.dvc Outputs data.xml prepare Outputs data/prepared featurize Outputs data/features @@ -22,7 +22,6 @@ describe('collectStages', () => { [join(dvcDemoPath, 'nested2')]: 'data/data.xml.dvc Outputs data.xml', [join(dvcDemoPath, 'nested3')]: undefined }) - expect(stages).toStrictEqual(['prepare', 'featurize', 'train', 'evaluate']) expect(pipelines).toStrictEqual( new Set([join(dvcDemoPath, 'nested1'), join(dvcDemoPath, 'nested3')]) diff --git a/extension/src/pipeline/model/collect.ts b/extension/src/pipeline/model/collect.ts index 52e3facae6..ee9aeb7c8b 100644 --- a/extension/src/pipeline/model/collect.ts +++ b/extension/src/pipeline/model/collect.ts @@ -1,31 +1,36 @@ import { trimAndSplit } from '../../util/stdout' -export const collectStages = (data: { - [pipeline: string]: string | undefined -}): { - pipelines: Set - stages: string[] - // eslint-disable-next-line sonarjs/cognitive-complexity -} => { - const stages: string[] = [] - const pipelines = new Set() +const failedValidation = (stageList: string | undefined): boolean => + stageList === undefined - for (const [pipeline, stageList] of Object.entries(data)) { - if (stageList === undefined) { - pipelines.add(pipeline) - } +const hasStages = (stageList: string | undefined): stageList is string => + !!stageList + +const hasValidStage = (stageList: string | undefined): boolean => { + if (!hasStages(stageList)) { + return false + } - if (!stageList) { + for (const stageStr of trimAndSplit(stageList)) { + const stage = stageStr.split(/\s+/)?.[0]?.trim() + if (!stage || stage.endsWith('.dvc')) { continue } - for (const stageStr of trimAndSplit(stageList)) { - const stage = stageStr.split(/\s+/)?.[0]?.trim() - if (!stage || stage.endsWith('.dvc')) { - continue - } - stages.push(stage) + + return true + } + return false +} + +export const collectPipelines = (data: { + [pipeline: string]: string | undefined +}): Set => { + const pipelines = new Set() + + for (const [pipeline, stageList] of Object.entries(data)) { + if (failedValidation(stageList) || hasValidStage(stageList)) { pipelines.add(pipeline) } } - return { pipelines, stages } + return pipelines } diff --git a/extension/src/pipeline/model/index.ts b/extension/src/pipeline/model/index.ts index 449091c864..a3ca71fd5b 100644 --- a/extension/src/pipeline/model/index.ts +++ b/extension/src/pipeline/model/index.ts @@ -1,36 +1,19 @@ import isEqual from 'lodash.isequal' -import { collectStages } from './collect' +import { collectPipelines } from './collect' import { Disposable } from '../../class/dispose' export class PipelineModel extends Disposable { - private readonly dvcRoot: string - private stages: string[] | undefined = [] private pipelines: Set | undefined - constructor(dvcRoot: string) { - super() - this.dvcRoot = dvcRoot - } - public transformAndSet(data: { [pipeline: string]: string | undefined }) { if (isEqual(data, {})) { - this.stages = undefined this.pipelines = undefined return } - const { pipelines, stages } = collectStages(data) + const pipelines = collectPipelines(data) this.pipelines = pipelines - this.stages = stages - } - - public hasStage() { - return !!(this.stages && this.stages.length > 0) - } - - public getStages() { - return this.stages } public getPipelines() { diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index 360af33b53..66847f4006 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/extension/src/test/suite/experiments/util.ts @@ -147,7 +147,7 @@ export const buildMultiRepoExperiments = (disposer: SafeWatcherDisposer) => { dvcRoot: dvcDemoPath, internalCommands }) - stub(pipeline, 'hasStage').returns(true) + stub(pipeline, 'hasPipeline').returns(true) const [experiments] = workspaceExperiments.create( [dvcDemoPath], diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts index 0809a86de5..7540a27357 100644 --- a/extension/src/test/suite/pipeline/index.test.ts +++ b/extension/src/test/suite/pipeline/index.test.ts @@ -41,41 +41,38 @@ suite('Pipeline Test Suite', () => { }) it('should return a cwd if there is an invalid pipeline (let command fail)', async () => { - const { pipeline, pipelineModel } = buildPipeline({ + const { pipeline } = buildPipeline({ disposer: disposable, dvcRoot: dvcDemoPath, stageList: null }) await pipeline.isReady() - expect(pipelineModel.hasStage()).to.be.false - expect(pipelineModel.hasPipeline()).to.be.true + expect(pipeline.hasPipeline()).to.be.true const cwd = await pipeline.getCwd() expect(cwd).to.equal(dvcDemoPath) }) it('should return a cwd if there is a single valid pipeline', async () => { - const { pipeline, pipelineModel } = buildPipeline({ + const { pipeline } = buildPipeline({ disposer: disposable, dvcRoot: dvcDemoPath, stageList: 'train' }) await pipeline.isReady() - expect(pipelineModel.hasStage()).to.be.true - expect(pipelineModel.hasPipeline()).to.be.true + expect(pipeline.hasPipeline()).to.be.true const cwd = await pipeline.getCwd() expect(cwd).to.equal(dvcDemoPath) }) it('should return the project root if there are multiple pipelines but one is the root', async () => { - const { pipeline, pipelineModel } = buildPipeline({ + const { pipeline } = buildPipeline({ disposer: disposable, dvcRoot: dvcDemoPath, dvcYamls: [dvcDemoPath, join(dvcDemoPath, 'subdir')], stageList: 'train' }) await pipeline.isReady() - expect(pipelineModel.hasStage()).to.be.true - expect(pipelineModel.hasPipeline()).to.be.true + expect(pipeline.hasPipeline()).to.be.true const cwd = await pipeline.getCwd() expect(cwd).to.equal(dvcDemoPath) }) @@ -87,15 +84,14 @@ suite('Pipeline Test Suite', () => { Thenable > mockShowQuickPick.resolves(pickedPipeline) - const { pipeline, pipelineModel } = buildPipeline({ + const { pipeline } = buildPipeline({ disposer: disposable, dvcRoot: dvcDemoPath, dvcYamls: [pickedPipeline, join(dvcDemoPath, 'nested2', 'dvc.yaml')], stageList: 'train' }) await pipeline.isReady() - expect(pipelineModel.hasStage()).to.be.true - expect(pipelineModel.hasPipeline()).to.be.true + expect(pipeline.hasPipeline()).to.be.true const cwd = await pipeline.getCwd() expect(cwd).to.equal(pickedPipeline) expect(mockShowQuickPick).to.be.calledOnce diff --git a/extension/src/test/suite/pipeline/util.ts b/extension/src/test/suite/pipeline/util.ts index dd5882eb46..5ade2f66d1 100644 --- a/extension/src/test/suite/pipeline/util.ts +++ b/extension/src/test/suite/pipeline/util.ts @@ -6,7 +6,6 @@ import { Pipeline } from '../../../pipeline' import { PipelineData } from '../../../pipeline/data' import { dvcDemoPath } from '../../util' import { buildDependencies } from '../util' -import { PipelineModel } from '../../../pipeline/model' export const buildExperimentsPipeline = ({ disposer, @@ -60,8 +59,6 @@ export const buildPipeline = ({ internalCommands, pipeline, // eslint-disable-next-line @typescript-eslint/no-explicit-any - pipelineData: (pipeline as any).data as PipelineData, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - pipelineModel: (pipeline as any).model as PipelineModel + pipelineData: (pipeline as any).data as PipelineData } } From 58034f6d1ba89924d948d51b3a0f157c868284fb Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Fri, 14 Jul 2023 10:35:23 +1000 Subject: [PATCH 11/12] refactor workspace experiments --- extension/src/experiments/index.ts | 15 +++++++-- extension/src/experiments/workspace.ts | 32 ++++++++++--------- .../src/test/suite/pipeline/index.test.ts | 2 +- 3 files changed, 30 insertions(+), 19 deletions(-) diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 68723023f4..30ef1a1a92 100644 --- a/extension/src/experiments/index.ts +++ b/extension/src/experiments/index.ts @@ -431,9 +431,13 @@ export class Experiments extends BaseRepository { } public async modifyWorkspaceParamsAndRun( - commandId: ModifiedExperimentAndRunCommandId, - cwd: string + commandId: ModifiedExperimentAndRunCommandId ) { + const cwd = await this.getPipelineCwd() + if (!cwd) { + return + } + const paramsToModify = await this.pickAndModifyParams() if (!paramsToModify) { return @@ -447,7 +451,12 @@ export class Experiments extends BaseRepository { return this.notifyChanged() } - public async modifyWorkspaceParamsAndQueue(cwd: string) { + public async modifyWorkspaceParamsAndQueue() { + const cwd = await this.getPipelineCwd() + if (!cwd) { + return + } + const paramsToModify = await this.pickAndModifyParams() if (!paramsToModify) { return diff --git a/extension/src/experiments/workspace.ts b/extension/src/experiments/workspace.ts index 6531c60c9e..ca1217d2dc 100644 --- a/extension/src/experiments/workspace.ts +++ b/extension/src/experiments/workspace.ts @@ -153,21 +153,29 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< commandId: ModifiedExperimentAndRunCommandId, overrideRoot?: string ) { - const { cwd, repository } = await this.getRepositoryAndCwd(overrideRoot) - if (!(cwd && repository)) { + const project = await this.getDvcRoot(overrideRoot) + if (!project) { + return + } + const repository = this.getRepository(project) + if (!repository) { return } - return await repository.modifyWorkspaceParamsAndRun(commandId, cwd) + return await repository.modifyWorkspaceParamsAndRun(commandId) } public async modifyWorkspaceParamsAndQueue(overrideRoot?: string) { - const { cwd, repository } = await this.getRepositoryAndCwd(overrideRoot) - if (!(cwd && repository)) { + const project = await this.getDvcRoot(overrideRoot) + if (!project) { + return + } + const repository = this.getRepository(project) + if (!repository) { return } - return await repository.modifyWorkspaceParamsAndQueue(cwd) + return await repository.modifyWorkspaceParamsAndQueue() } public async getCwdThenRun(commandId: CommandId) { @@ -407,19 +415,13 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews< return this.getRepository(dvcRoot)[method]() } - private async getRepositoryAndCwd(overrideRoot?: string) { + private async getCwd(overrideRoot?: string) { const project = await this.getDvcRoot(overrideRoot) if (!project) { - return { cwd: undefined, repository: undefined } + return } const repository = this.getRepository(project) - const cwd = await repository.getPipelineCwd() - return { cwd, repository } - } - - private async getCwd(overrideRoot?: string) { - const { cwd } = await this.getRepositoryAndCwd(overrideRoot) - return cwd + return await repository.getPipelineCwd() } private async pickExpThenRun( diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts index 7540a27357..5d97364a24 100644 --- a/extension/src/test/suite/pipeline/index.test.ts +++ b/extension/src/test/suite/pipeline/index.test.ts @@ -161,7 +161,7 @@ suite('Pipeline Test Suite', () => { ) }) - it('should ask to enter a custom command if the file is not a python file or Jupyter notebook', async () => { + it('should ask to enter a custom command if the file is not a Python file or Jupyter notebook', async () => { const mockNow = getMockNow() const { pipeline } = buildPipeline({ disposer: disposable, From 2c77241c180d23fb5fc1a9d365c50f59bfdbc9a3 Mon Sep 17 00:00:00 2001 From: Matt Seddon Date: Fri, 14 Jul 2023 10:41:32 +1000 Subject: [PATCH 12/12] add more unit tests for pipeline collection --- extension/src/pipeline/model/collect.test.ts | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/extension/src/pipeline/model/collect.test.ts b/extension/src/pipeline/model/collect.test.ts index d6520cb1c7..88439fa8de 100644 --- a/extension/src/pipeline/model/collect.test.ts +++ b/extension/src/pipeline/model/collect.test.ts @@ -27,4 +27,20 @@ describe('collectPipelines', () => { new Set([join(dvcDemoPath, 'nested1'), join(dvcDemoPath, 'nested3')]) ) }) + + it('should not collect a pipeline that has no stages', () => { + const pipelines = collectPipelines({ + [dvcDemoPath]: '' + }) + + expect(pipelines).toStrictEqual(new Set([])) + }) + + it('should collect a pipeline that failed validation', () => { + const pipelines = collectPipelines({ + [dvcDemoPath]: undefined + }) + + expect(pipelines).toStrictEqual(new Set([dvcDemoPath])) + }) })