diff --git a/extension/src/experiments/data/index.ts b/extension/src/experiments/data/index.ts index 71c76de00c..9f4149367c 100644 --- a/extension/src/experiments/data/index.ts +++ b/extension/src/experiments/data/index.ts @@ -7,7 +7,7 @@ import { import { getRelativePattern } from '../../fileSystem/relativePattern' import { createFileSystemWatcher } from '../../fileSystem/watcher' import { AvailableCommands, InternalCommands } from '../../commands/internal' -import { ExpShowOutput } from '../../cli/dvc/contract' +import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../../cli/dvc/contract' import { BaseData } from '../../data' import { DOT_DVC, ExperimentFlag } from '../../cli/dvc/constants' import { gitPath } from '../../cli/git/constants' @@ -47,33 +47,36 @@ export class ExperimentsData extends BaseData { ) void this.updateAvailableBranchesToSelect(allBranches) + const data: ExpShowOutput = [] const { branches, currentBranch } = await this.getBranchesToShowWithCurrent( allBranches ) - const flags = [] - const branchList: string[] = [] - - for (const branch of branches) { - const nbOfCommitsToShow = this.experiments.getNbOfCommitsToShow(branch) - for (let i = 0; i < nbOfCommitsToShow; i++) { - flags.push(ExperimentFlag.REV, `${branch}~${i}`) - branchList.push(branch) - } - } - - const output = await this.internalCommands.executeCommand( - AvailableCommands.EXP_SHOW, - this.dvcRoot, - ...flags + await Promise.all( + branches.map(async branch => { + const branchFlags = [ + ExperimentFlag.REV, + branch, + ExperimentFlag.NUM_COMMIT, + this.experiments.getNbOfCommitsToShow(branch).toString() + ] + + const output = (await this.expShow( + branchFlags, + branch + )) as ExpShowOutput + + if (branch !== currentBranch) { + const workspaceIndex = output.findIndex( + exp => exp.rev === EXPERIMENT_WORKSPACE_ID + ) + output.splice(workspaceIndex, 1) + } + data.push(...output) + }) ) - const data = output.map((out, i) => ({ - ...out, - branch: i === 0 ? currentBranch : branchList[i - 1] - })) - this.collectFiles(data) return this.notifyChanged(data) @@ -97,10 +100,18 @@ export class ExperimentsData extends BaseData { branches.push(currentBranch) this.experiments.setBranchesToShow(branches) } - return { branches, currentBranch } } + private async expShow(flags: (ExperimentFlag | string)[], branch?: string) { + const data = await this.internalCommands.executeCommand( + AvailableCommands.EXP_SHOW, + this.dvcRoot, + ...flags + ) + return data.map(exp => ({ ...exp, branch })) + } + private async updateAvailableBranchesToSelect(branches?: string[]) { const allBranches = branches || diff --git a/extension/src/test/suite/experiments/data/index.test.ts b/extension/src/test/suite/experiments/data/index.test.ts index b68dc6edbd..63f7791f76 100644 --- a/extension/src/test/suite/experiments/data/index.test.ts +++ b/extension/src/test/suite/experiments/data/index.test.ts @@ -19,10 +19,7 @@ import { InternalCommands } from '../../../../commands/internal' import { buildExperimentsData } from '../util' -import { - DEFAULT_NUM_OF_COMMITS_TO_SHOW, - ExperimentFlag -} from '../../../../cli/dvc/constants' +import { DEFAULT_NUM_OF_COMMITS_TO_SHOW } from '../../../../cli/dvc/constants' import { EXPERIMENTS_GIT_LOGS_REFS } from '../../../../experiments/data/constants' import { gitPath } from '../../../../cli/git/constants' import * as FileSystem from '../../../../fileSystem' @@ -187,6 +184,24 @@ suite('Experiments Data Test Suite', () => { expect(managedUpdateSpy).to.be.called }).timeout(10000) + it('should call exp show for each branch to show', async () => { + stub(ExperimentsData.prototype, 'managedUpdate').resolves() + const branchesToShow = [ + 'main', + 'my-other-branch', + 'secret-branch', + 'old-branch' + ] + const { data, mockExpShow, mockGetBranchesToShow } = + buildExperimentsData(disposable) + + mockGetBranchesToShow.returns(branchesToShow) + + await data.update() + + expect(mockExpShow).to.have.callCount(branchesToShow.length) + }) + it('should prune any old branches to show before calling exp show on them', async () => { stub(ExperimentsData.prototype, 'managedUpdate').resolves() const branchesToShow = [ @@ -211,19 +226,11 @@ suite('Experiments Data Test Suite', () => { 'branch-283498' ) - mockGetBranchesToShow.returns([]) + mockGetBranchesToShow.returns(['main']) await data.update() - expect(mockExpShow).to.have.been.calledWithMatch( - dvcDemoPath, - ExperimentFlag.REV, - 'branch-283498~0', - ExperimentFlag.REV, - 'branch-283498~1', - ExperimentFlag.REV, - 'branch-283498~2' - ) + expect(mockExpShow).to.have.callCount(2) }) }) })