diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 8e1f100457..386169b69f 100644 --- a/extension/src/experiments/index.ts +++ b/extension/src/experiments/index.ts @@ -107,8 +107,8 @@ export class Experiments { return this.getRepository(dvcRoot).getFilter(id) } - public getRunningOrQueued(dvcRoot: string): string[] { - return this.getRepository(dvcRoot).getRunningOrQueued() + public getExperimentNames(dvcRoot: string): string[] { + return this.getRepository(dvcRoot).getExperimentNames() } public getExperiment(dvcRoot: string, name: string) { diff --git a/extension/src/experiments/model.ts b/extension/src/experiments/model.ts index eea9f03af0..11676ad2ff 100644 --- a/extension/src/experiments/model.ts +++ b/extension/src/experiments/model.ts @@ -114,10 +114,11 @@ export class ExperimentsModel { return this.columnStatus[path] } - public getRunningOrQueued(): string[] { - return [this.workspace, ...this.getExperiments()] - .filter(experiment => experiment?.queued || experiment?.running) - .map(experiment => experiment.displayName) + public getExperimentNames(): string[] { + const workspace = this.workspace.running ? [this.workspace] : [] + return [...workspace, ...this.getExperiments()].map( + experiment => experiment?.displayName + ) } public getExperiment(name: string) { diff --git a/extension/src/experiments/repository.ts b/extension/src/experiments/repository.ts index 81d686c38b..0d97d2f166 100644 --- a/extension/src/experiments/repository.ts +++ b/extension/src/experiments/repository.ts @@ -175,8 +175,8 @@ export class ExperimentsRepository { } } - public getRunningOrQueued(): string[] { - return this.model.getRunningOrQueued() + public getExperimentNames(): string[] { + return this.model.getExperimentNames() } public getExperiment(name: string) { diff --git a/extension/src/experiments/views/runsTree.test.ts b/extension/src/experiments/views/runsTree.test.ts index f14f2136c1..dd14e45f43 100644 --- a/extension/src/experiments/views/runsTree.test.ts +++ b/extension/src/experiments/views/runsTree.test.ts @@ -20,13 +20,13 @@ const mockedDisposable = mocked(Disposable) const mockedGetDvcRoots = jest.fn() const mockedGetExperiment = jest.fn() const mockedGetCheckpointNames = jest.fn() -const mockedGetRunningOrQueued = jest.fn() +const mockedGetExperimentNames = jest.fn() const mockedExperiments = { experimentsRowsChanged: mockedExperimentsRowsChanged, getCheckpointNames: mockedGetCheckpointNames, getDvcRoots: mockedGetDvcRoots, getExperiment: mockedGetExperiment, - getRunningOrQueued: mockedGetRunningOrQueued, + getExperimentNames: mockedGetExperimentNames, isReady: () => true } as unknown as Experiments @@ -48,7 +48,7 @@ describe('ExperimentsRunsTree', () => { it('should return an empty array when no experiments are queued for any of the repositories', async () => { const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) mockedGetDvcRoots.mockReturnValueOnce(['demo']) - mockedGetRunningOrQueued.mockReturnValueOnce([]) + mockedGetExperimentNames.mockReturnValueOnce([]) const rootElements = await experimentsRunsTree.getChildren() @@ -59,9 +59,9 @@ describe('ExperimentsRunsTree', () => { const dvcRoots = ['demo', 'and/mock', 'other/repo'] const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) mockedGetDvcRoots.mockReturnValueOnce(dvcRoots) - mockedGetRunningOrQueued.mockReturnValueOnce([]) - mockedGetRunningOrQueued.mockReturnValueOnce([]) - mockedGetRunningOrQueued.mockReturnValueOnce(['90aea7f']) + mockedGetExperimentNames.mockReturnValueOnce([]) + mockedGetExperimentNames.mockReturnValueOnce([]) + mockedGetExperimentNames.mockReturnValueOnce(['90aea7f']) const rootElements = await experimentsRunsTree.getChildren() @@ -71,8 +71,8 @@ describe('ExperimentsRunsTree', () => { it('should return an array of queued experiment names when an element is provided', async () => { const queuedExperiments = ['90aea7f', 'f0778b3', 'f81f1b5'] const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) - mockedGetRunningOrQueued.mockReturnValueOnce(queuedExperiments) - mockedGetRunningOrQueued.mockReturnValueOnce(queuedExperiments) + mockedGetExperimentNames.mockReturnValueOnce(queuedExperiments) + mockedGetExperimentNames.mockReturnValueOnce(queuedExperiments) mockedGetDvcRoots.mockReturnValueOnce(['repo']) await experimentsRunsTree.getChildren() @@ -83,15 +83,10 @@ describe('ExperimentsRunsTree', () => { }) it('should return an array of checkpoints when a non root element is provided', async () => { - const runningOrQueuedExperiments = [ - 'ebbd66f', - 'e5855d7', - '6e5e782', - '15c9c56' - ] + const experimentNames = ['ebbd66f', 'e5855d7', '6e5e782', '15c9c56'] const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) - mockedGetRunningOrQueued.mockReturnValueOnce(runningOrQueuedExperiments) - mockedGetRunningOrQueued.mockReturnValueOnce(runningOrQueuedExperiments) + mockedGetExperimentNames.mockReturnValueOnce(experimentNames) + mockedGetExperimentNames.mockReturnValueOnce(experimentNames) mockedGetDvcRoots.mockReturnValueOnce(['repo']) await experimentsRunsTree.getChildren() @@ -117,7 +112,7 @@ describe('ExperimentsRunsTree', () => { const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) mockedGetDvcRoots.mockReturnValueOnce(['demo']) - mockedGetRunningOrQueued.mockReturnValueOnce([]) + mockedGetExperimentNames.mockReturnValueOnce([]) await experimentsRunsTree.getChildren() @@ -139,8 +134,8 @@ describe('ExperimentsRunsTree', () => { const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) mockedGetDvcRoots.mockReturnValueOnce(['demo']) const mockedQueuedExperiment = 'f0778b3' - mockedGetRunningOrQueued.mockReturnValueOnce([mockedQueuedExperiment]) - mockedGetRunningOrQueued.mockReturnValueOnce([mockedQueuedExperiment]) + mockedGetExperimentNames.mockReturnValueOnce([mockedQueuedExperiment]) + mockedGetExperimentNames.mockReturnValueOnce([mockedQueuedExperiment]) mockedGetExperiment.mockReturnValueOnce({ queued: true, running: false @@ -168,8 +163,8 @@ describe('ExperimentsRunsTree', () => { const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) mockedGetDvcRoots.mockReturnValueOnce(['demo']) const mockedRunningExperiment = 'workspace' - mockedGetRunningOrQueued.mockReturnValueOnce([mockedRunningExperiment]) - mockedGetRunningOrQueued.mockReturnValueOnce([mockedRunningExperiment]) + mockedGetExperimentNames.mockReturnValueOnce([mockedRunningExperiment]) + mockedGetExperimentNames.mockReturnValueOnce([mockedRunningExperiment]) await experimentsRunsTree.getChildren() await experimentsRunsTree.getChildren('demo') @@ -195,8 +190,8 @@ describe('ExperimentsRunsTree', () => { const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) mockedGetDvcRoots.mockReturnValueOnce(['demo']) const mockedRunningExperiment = 'f0778b3' - mockedGetRunningOrQueued.mockReturnValueOnce([mockedRunningExperiment]) - mockedGetRunningOrQueued.mockReturnValueOnce([mockedRunningExperiment]) + mockedGetExperimentNames.mockReturnValueOnce([mockedRunningExperiment]) + mockedGetExperimentNames.mockReturnValueOnce([mockedRunningExperiment]) mockedGetExperiment.mockReturnValueOnce({ hasChildren: true, running: true @@ -226,6 +221,33 @@ describe('ExperimentsRunsTree', () => { const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) const treeItem = experimentsRunsTree.getTreeItem('f0778b3') + expect(treeItem).toEqual({ + ...mockedItem, + iconPath: { id: 'debug-stackframe-dot' } + }) + }) + + it('should return a tree item for an existing experiment', async () => { + let mockedItem = {} + mockedTreeItem.mockImplementationOnce(function (uri, collapsibleState) { + mockedItem = { collapsibleState, uri } + return mockedItem + }) + mockedThemeIcon.mockImplementationOnce(function (id) { + return { id } + }) + + const experimentsRunsTree = new ExperimentsRunsTree(mockedExperiments) + mockedGetDvcRoots.mockReturnValueOnce(['demo']) + const mockedExistingExperiment = 'f0998a3' + mockedGetExperimentNames.mockReturnValueOnce([mockedExistingExperiment]) + mockedGetExperimentNames.mockReturnValueOnce([mockedExistingExperiment]) + mockedGetExperiment.mockReturnValueOnce({}) + + await experimentsRunsTree.getChildren() + await experimentsRunsTree.getChildren('demo') + + const treeItem = experimentsRunsTree.getTreeItem(mockedExistingExperiment) expect(treeItem).toEqual({ ...mockedItem, iconPath: { id: 'primitive-dot' } diff --git a/extension/src/experiments/views/runsTree.ts b/extension/src/experiments/views/runsTree.ts index 97653eef52..f26765d972 100644 --- a/extension/src/experiments/views/runsTree.ts +++ b/extension/src/experiments/views/runsTree.ts @@ -38,18 +38,7 @@ export class ExperimentsRunsTree implements TreeDataProvider { return new TreeItem(Uri.file(element), TreeItemCollapsibleState.Collapsed) } - const dvcRoot = this.runRoots[element] - if (!dvcRoot) { - return this.getRunningCheckpoint(element) - } - - const experiment = this.experiments.getExperiment(dvcRoot, element) - - if (element === 'workspace' || experiment?.running) { - return this.getRunning(element, experiment?.hasChildren) - } - - return this.getQueued(element) + return this.getExperimentTreeItem(element) } public getChildren(element?: string): Promise { @@ -58,7 +47,7 @@ export class ExperimentsRunsTree implements TreeDataProvider { } if (this.isRoot(element)) { - return Promise.resolve(this.getRunningOrQueued(element)) + return Promise.resolve(this.getExperimentNames(element)) } return Promise.resolve( @@ -66,11 +55,30 @@ export class ExperimentsRunsTree implements TreeDataProvider { ) } + private getExperimentTreeItem(element: string) { + const dvcRoot = this.runRoots[element] + if (!dvcRoot) { + return this.getRunningCheckpoint(element) + } + + const experiment = this.experiments.getExperiment(dvcRoot, element) + + if (element === 'workspace' || experiment?.running) { + return this.getRunning(element, experiment?.hasChildren) + } + + if (experiment?.queued) { + return this.getQueued(element) + } + + return this.getExistingExperiment(element, experiment?.hasChildren) + } + private getRunningCheckpoint(element: string) { return this.getTreeItemWithIcon( element, TreeItemCollapsibleState.None, - 'primitive-dot' + 'debug-stackframe-dot' ) } @@ -90,6 +98,16 @@ export class ExperimentsRunsTree implements TreeDataProvider { ) } + private getExistingExperiment(element: string, hasChildren?: boolean) { + return this.getTreeItemWithIcon( + element, + hasChildren + ? TreeItemCollapsibleState.Collapsed + : TreeItemCollapsibleState.None, + 'primitive-dot' + ) + } + private getTreeItemWithIcon( element: string, collapsibleState: TreeItemCollapsibleState, @@ -103,23 +121,23 @@ export class ExperimentsRunsTree implements TreeDataProvider { private async getRootElements() { await this.experiments.isReady() const dvcRoots = this.experiments.getDvcRoots() - const runningOrQueued = flatten( + const experimentNames = flatten( dvcRoots.map(dvcRoot => { this.runRoots[dvcRoot] = dvcRoot - return this.experiments.getRunningOrQueued(dvcRoot) + return this.experiments.getExperimentNames(dvcRoot) }) ) - if (definedAndNonEmpty(runningOrQueued)) { + if (definedAndNonEmpty(experimentNames)) { return dvcRoots.sort((a, b) => a.localeCompare(b)) } return [] } - private getRunningOrQueued(dvcRoot: string): string[] { - const runningOrQueued = this.experiments.getRunningOrQueued(dvcRoot) - runningOrQueued.forEach(name => (this.runRoots[name] = dvcRoot)) - return runningOrQueued + private getExperimentNames(dvcRoot: string): string[] { + const experimentNames = this.experiments.getExperimentNames(dvcRoot) + experimentNames.forEach(name => (this.runRoots[name] = dvcRoot)) + return experimentNames } private isRoot(element: string) { diff --git a/extension/src/test/suite/experiments/repository.test.ts b/extension/src/test/suite/experiments/repository.test.ts index a6a0430c5b..ee789764f3 100644 --- a/extension/src/test/suite/experiments/repository.test.ts +++ b/extension/src/test/suite/experiments/repository.test.ts @@ -68,8 +68,8 @@ suite('Experiments Repository Test Suite', () => { }) }) - describe('getRunningOrQueued', () => { - it('should return the currently running or queued experiments', async () => { + describe('getExperimentNames', () => { + it("should return all existing experiments' names", async () => { const config = disposable.track(new Config()) const cliReader = disposable.track(new CliReader(config)) stub(cliReader, 'experimentShow').resolves(complexExperimentsOutput) @@ -87,11 +87,14 @@ suite('Experiments Repository Test Suite', () => { ) await experimentsRepository.isReady() - const runningOrQueued = experimentsRepository.getRunningOrQueued() + const experimentNames = experimentsRepository.getExperimentNames() - expect(runningOrQueued).to.deep.equal([ + expect(experimentNames).to.deep.equal([ 'workspace', + 'exp-05694', 'exp-e7a67', + 'test-branch', + 'exp-83425', '90aea7f' ]) })