Cleanup old branches before showing the experiments table (#3917)

* Cleanup old branches before showing

* Fix and add tests

Author: sroy3

extension/src/experiments/data/index.ts

@@ -41,11 +41,17 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
   }
 
   public async update(): Promise<void> {
-    void this.updateAvailableBranchesToSelect()
+    const allBranches = await this.internalCommands.executeCommand<string[]>(
+      AvailableCommands.GIT_GET_BRANCHES,
+      this.dvcRoot
+    )
+
+    void this.updateAvailableBranchesToSelect(allBranches)
     const data: ExpShowOutput = []
 
-    const { branches, currentBranch } =
-      await this.getBranchesToShowWithCurrent()
+    const { branches, currentBranch } = await this.getBranchesToShowWithCurrent(
+      allBranches
+    )
 
     await Promise.all(
       branches.map(async branch => {
@@ -55,6 +61,7 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
           ExperimentFlag.NUM_COMMIT,
           this.experiments.getNbOfCommitsToShow(branch).toString()
         ]
+
         const output = (await this.expShow(
           branchFlags,
           branch
@@ -79,12 +86,14 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
     this.collectedFiles = collectFiles(data, this.collectedFiles)
   }
 
-  private async getBranchesToShowWithCurrent() {
+  private async getBranchesToShowWithCurrent(allBranches: string[]) {
     const currentBranch = await this.internalCommands.executeCommand<string>(
       AvailableCommands.GIT_GET_CURRENT_BRANCH,
       this.dvcRoot
     )
 
+    this.experiments.pruneBranchesToShow(allBranches)
+
     const branches = this.experiments.getBranchesToShow()
 
     if (!branches.includes(currentBranch)) {
@@ -103,11 +112,13 @@ export class ExperimentsData extends BaseData<ExpShowOutput> {
     return data.map(exp => ({ ...exp, branch }))
   }
 
-  private async updateAvailableBranchesToSelect() {
-    const allBranches = await this.internalCommands.executeCommand<string[]>(
-      AvailableCommands.GIT_GET_BRANCHES,
-      this.dvcRoot
-    )
+  private async updateAvailableBranchesToSelect(branches?: string[]) {
+    const allBranches =
+      branches ||
+      (await this.internalCommands.executeCommand<string[]>(
+        AvailableCommands.GIT_GET_BRANCHES,
+        this.dvcRoot
+      ))
     this.experiments.setAvailableBranchesToShow(allBranches)
   }
 

extension/src/experiments/model/index.test.ts

@@ -335,4 +335,27 @@ describe('ExperimentsModel', () => {
       }
     )
   })
+
+  it('should remove outdated branches to show when calling pruneBranchesToShow', () => {
+    const model = new ExperimentsModel('', buildMockMemento())
+
+    model.setBranchesToShow(['one', 'old', 'two', 'three', 'older'])
+    model.pruneBranchesToShow(['one', 'two', 'three', 'four', 'five', 'six'])
+
+    expect(model.getBranchesToShow()).toStrictEqual(['one', 'two', 'three'])
+  })
+
+  it('should persist the branches to show when calling pruneBranchesToShow', () => {
+    const memento = buildMockMemento()
+    const model = new ExperimentsModel('', memento)
+
+    model.setBranchesToShow(['one', 'old', 'two', 'three', 'older'])
+    model.pruneBranchesToShow(['one', 'two', 'three', 'four', 'five', 'six'])
+
+    expect(memento.get(PersistenceKey.EXPERIMENTS_BRANCHES)).toStrictEqual([
+      'one',
+      'two',
+      'three'
+    ])
+  })
 })

extension/src/experiments/model/index.ts

@@ -429,6 +429,13 @@ export class ExperimentsModel extends ModelWithPersistence {
     this.persistBranchesToShow()
   }
 
+  public pruneBranchesToShow(branches: string[]) {
+    this.branchesToShow = this.branchesToShow.filter(branch =>
+      branches.includes(branch)
+    )
+    this.persistBranchesToShow()
+  }
+
   public getBranchesToShow() {
     return this.branchesToShow
   }

extension/src/test/suite/experiments/data/index.test.ts

@@ -47,6 +47,7 @@ suite('Experiments Data Test Suite', () => {
 
     it('should debounce all calls to update that are made within 200ms', async () => {
       const { data, mockExpShow } = buildExperimentsData(disposable)
+      await data.isReady()
 
       await Promise.all([
         data.managedUpdate(),
@@ -97,6 +98,7 @@ suite('Experiments Data Test Suite', () => {
             getNbOfCommitsToShow: () => ({
               main: DEFAULT_NUM_OF_COMMITS_TO_SHOW
             }),
+            pruneBranchesToShow: stub(),
             setAvailableBranchesToShow: stub(),
             setBranchesToShow: stub()
           } as unknown as ExperimentsModel
@@ -156,6 +158,7 @@ suite('Experiments Data Test Suite', () => {
             getNbOfCommitsToShow: () => ({
               main: DEFAULT_NUM_OF_COMMITS_TO_SHOW
             }),
+            pruneBranchesToShow: stub(),
             setAvailableBranchesToShow: stub(),
             setBranchesToShow: stub()
           } as unknown as ExperimentsModel
@@ -199,6 +202,23 @@ suite('Experiments Data Test Suite', () => {
       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 = [
+        'main',
+        'my-other-branch',
+        'secret-branch',
+        'old-branch'
+      ]
+      const { data, mockPruneBranchesToShow, mockGetBranchesToShow } =
+        buildExperimentsData(disposable)
+      mockGetBranchesToShow.returns(branchesToShow)
+
+      await data.update()
+
+      expect(mockPruneBranchesToShow).to.be.calledOnce
+    })
+
     it('should add the current branch to the exp show output', async () => {
       stub(ExperimentsData.prototype, 'managedUpdate').resolves()
       const { data, mockExpShow, mockGetBranchesToShow } = buildExperimentsData(

extension/src/test/suite/experiments/util.ts

@@ -16,7 +16,6 @@ import { ExperimentsModel } from '../../../experiments/model'
 import { ColumnsModel } from '../../../experiments/columns/model'
 import { DEFAULT_NUM_OF_COMMITS_TO_SHOW } from '../../../cli/dvc/constants'
 import { PersistenceKey } from '../../../persistence/constants'
-import { GitReader } from '../../../cli/git/reader'
 
 export const buildExperiments = (
   disposer: Disposer,
@@ -142,27 +141,38 @@ const buildExperimentsDataDependencies = (disposer: Disposer) => {
     'createFileSystemWatcher'
   ).returns(undefined)
 
-  const { dvcReader, internalCommands } = buildInternalCommands(disposer)
+  const { dvcReader, gitReader, internalCommands } =
+    buildInternalCommands(disposer)
   const mockExpShow = stub(dvcReader, 'expShow').resolves(expShowFixture)
-  return { internalCommands, mockCreateFileSystemWatcher, mockExpShow }
+  return {
+    gitReader,
+    internalCommands,
+    mockCreateFileSystemWatcher,
+    mockExpShow
+  }
 }
 
 export const buildExperimentsData = (
   disposer: SafeWatcherDisposer,
   currentBranch = 'main'
 ) => {
-  stub(GitReader.prototype, 'getCurrentBranch').resolves(currentBranch)
+  const {
+    internalCommands,
+    mockExpShow,
+    mockCreateFileSystemWatcher,
+    gitReader
+  } = buildExperimentsDataDependencies(disposer)
 
-  const { internalCommands, mockExpShow, mockCreateFileSystemWatcher } =
-    buildExperimentsDataDependencies(disposer)
+  stub(gitReader, 'getCurrentBranch').resolves(currentBranch)
+  stub(gitReader, 'getBranches').resolves(['one'])
 
   const mockGetBranchesToShow = stub().returns(['main'])
+  const mockPruneBranchesToShow = stub()
   const data = disposer.track(
     new ExperimentsData(dvcDemoPath, internalCommands, {
       getBranchesToShow: mockGetBranchesToShow,
-      getNbOfCommitsToShow: () => ({
-        main: DEFAULT_NUM_OF_COMMITS_TO_SHOW
-      }),
+      getNbOfCommitsToShow: () => DEFAULT_NUM_OF_COMMITS_TO_SHOW,
+      pruneBranchesToShow: mockPruneBranchesToShow,
       setAvailableBranchesToShow: stub(),
       setBranchesToShow: stub()
     } as unknown as ExperimentsModel)
@@ -172,7 +182,8 @@ export const buildExperimentsData = (
     data,
     mockCreateFileSystemWatcher,
     mockExpShow,
-    mockGetBranchesToShow
+    mockGetBranchesToShow,
+    mockPruneBranchesToShow
   }
 }
 

extension/src/test/suite/util.ts

@@ -188,6 +188,7 @@ export const buildMockExperimentsData = (update = stub()) =>
     getCurrentBranchesToShow: () => ['current'],
     onDidChangeDvcYaml: stub(),
     onDidUpdate: stub(),
+    pruneBranchesToShow: stub(),
     setCurrentBranchesToShow: stub(),
     update
   } as unknown as ExperimentsData)