Prevent current branch being duplicated when an experiment is running (#4084)

Author: mattseddon

extension/src/experiments/data/index.ts

@@ -33,22 +33,15 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
 
     void this.watchExpGitRefs()
     void this.managedUpdate()
-    void this.updateAvailableBranchesToSelect()
   }
 
   public managedUpdate() {
     return this.processManager.run('update')
   }
 
   public async update(): Promise<void> {
-    const allBranches = await this.internalCommands.executeCommand<string[]>(
-      AvailableCommands.GIT_GET_BRANCHES,
-      this.dvcRoot
-    )
-
-    void this.updateAvailableBranchesToSelect(allBranches)
-
-    const branches = await this.getBranchesToShow(allBranches)
+    await this.updateBranches()
+    const branches = this.experiments.getBranchesToShow()
     let gitLog = ''
     const rowOrder: { branch: string; sha: string }[] = []
     const availableNbCommits: { [branch: string]: number } = {}
@@ -115,34 +108,19 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     return gitLog
   }
 
-  private async getBranchesToShow(allBranches: string[]) {
-    const currentBranch = await this.internalCommands.executeCommand<string>(
-      AvailableCommands.GIT_GET_CURRENT_BRANCH,
-      this.dvcRoot
-    )
-
-    this.experiments.pruneBranchesToShow(allBranches)
-
-    const currentBranches = [
-      currentBranch,
-      ...this.experiments
-        .getBranchesToShow()
-        .filter(branch => branch !== currentBranch)
-    ]
-
-    this.experiments.setBranchesToShow(currentBranches)
-
-    return currentBranches
-  }
-
-  private async updateAvailableBranchesToSelect(branches?: string[]) {
-    const allBranches =
-      branches ||
-      (await this.internalCommands.executeCommand<string[]>(
+  private async updateBranches() {
+    const [currentBranch, allBranches] = await Promise.all([
+      this.internalCommands.executeCommand<string>(
+        AvailableCommands.GIT_GET_CURRENT_BRANCH,
+        this.dvcRoot
+      ),
+      this.internalCommands.executeCommand<string[]>(
         AvailableCommands.GIT_GET_BRANCHES,
         this.dvcRoot
-      ))
-    this.experiments.setAvailableBranchesToShow(allBranches)
+      )
+    ])
+
+    this.experiments.setBranches(currentBranch, allBranches)
   }
 
   private async watchExpGitRefs(): Promise<void> {

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

@@ -388,21 +388,30 @@ describe('ExperimentsModel', () => {
     )
   })
 
-  it('should remove outdated branches to show when calling pruneBranchesToShow', () => {
+  it('should always return the current branch first', () => {
     const model = new ExperimentsModel('', buildMockMemento())
 
-    model.setBranchesToShow(['A', 'old', 'B', 'C', 'older'])
-    model.pruneBranchesToShow(['A', 'B', 'C', 'four', 'five', 'six'])
+    model.setSelectedBranches(['A', 'old', 'B', 'C', 'older'])
+    model.setBranches('six', ['A', 'B', 'C', 'four', 'five', 'six'])
 
-    expect(model.getBranchesToShow()).toStrictEqual(['A', 'B', 'C'])
+    expect(model.getBranchesToShow()).toStrictEqual(['six', 'A', 'B', 'C'])
   })
 
-  it('should persist the branches to show when calling pruneBranchesToShow', () => {
+  it('should remove outdated selected branches when calling setBranches', () => {
+    const model = new ExperimentsModel('', buildMockMemento())
+
+    model.setSelectedBranches(['A', 'old', 'B', 'C', 'older'])
+    model.setBranches('main', ['A', 'B', 'C', 'four', 'five', 'six'])
+
+    expect(model.getBranchesToShow()).toStrictEqual(['main', 'A', 'B', 'C'])
+  })
+
+  it('should persist the selected branches when calling setBranches', () => {
     const memento = buildMockMemento()
     const model = new ExperimentsModel('', memento)
 
-    model.setBranchesToShow(['A', 'old', 'B', 'C', 'older'])
-    model.pruneBranchesToShow(['A', 'B', 'C', 'four', 'five', 'six'])
+    model.setSelectedBranches(['A', 'old', 'B', 'C', 'older'])
+    model.setBranches('main', ['A', 'B', 'C', 'four', 'five', 'six'])
 
     expect(memento.get(PersistenceKey.EXPERIMENTS_BRANCHES)).toStrictEqual([
       'A',
@@ -411,6 +420,35 @@ describe('ExperimentsModel', () => {
     ])
   })
 
+  it('should persist the selected branches when calling setSelectedBranches', () => {
+    const memento = buildMockMemento()
+    const model = new ExperimentsModel('', memento)
+
+    model.setSelectedBranches(['A', 'old', 'B', 'C', 'older'])
+
+    expect(memento.get(PersistenceKey.EXPERIMENTS_BRANCHES)).toStrictEqual([
+      'A',
+      'old',
+      'B',
+      'C',
+      'older'
+    ])
+  })
+
+  it('should remove the current branch from the selected branches and only persist the selected branches', () => {
+    const memento = buildMockMemento()
+    const model = new ExperimentsModel('', memento)
+
+    model.setSelectedBranches(['A', 'old', 'B', 'C', 'older'])
+    model.setBranches('A', ['A', 'B', 'C', 'four', 'five', 'six'])
+
+    expect(model.getBranchesToShow()).toStrictEqual(['A', 'B', 'C'])
+    expect(memento.get(PersistenceKey.EXPERIMENTS_BRANCHES)).toStrictEqual([
+      'B',
+      'C'
+    ])
+  })
+
   const getSelectedRevisions = (
     model: ExperimentsModel
   ): { id: string; displayColor: string }[] =>

extension/src/experiments/model/index.ts

@@ -62,7 +62,8 @@ export class ExperimentsModel extends ModelWithPersistence {
   private coloredStatus: ColoredStatus
   private starredExperiments: StarredExperiments
   private numberOfCommitsToShow: Record<string, number>
-  private branchesToShow: string[] = []
+  private currentBranch: string | undefined
+  private selectedBranches: string[] = []
   private availableBranchesToShow: string[] = []
   private hasMoreCommits: { [branch: string]: boolean } = {}
   private isShowingMoreCommits: { [branch: string]: boolean } = {}
@@ -102,7 +103,7 @@ export class ExperimentsModel extends ModelWithPersistence {
       this.numberOfCommitsToShow = {}
     }
 
-    this.branchesToShow = this.revive<string[]>(
+    this.selectedBranches = this.revive<string[]>(
       PersistenceKey.EXPERIMENTS_BRANCHES,
       []
     )
@@ -463,24 +464,24 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.numberOfCommitsToShow
   }
 
-  public setBranchesToShow(branches: string[]) {
-    this.branchesToShow = branches
+  public setBranches(currentBranch: string, allBranches: string[]) {
+    this.availableBranchesToShow = allBranches
+    this.currentBranch = currentBranch
+    this.selectedBranches = this.selectedBranches.filter(
+      branch => allBranches.includes(branch) && branch !== this.currentBranch
+    )
     this.persistBranchesToShow()
   }
 
-  public pruneBranchesToShow(branches: string[]) {
-    this.branchesToShow = this.branchesToShow.filter(branch =>
-      branches.includes(branch)
+  public setSelectedBranches(branches: string[]) {
+    this.selectedBranches = branches.filter(
+      branch => branch !== this.currentBranch
     )
     this.persistBranchesToShow()
   }
 
   public getBranchesToShow() {
-    return this.branchesToShow
-  }
-
-  public setAvailableBranchesToShow(branches: string[]) {
-    this.availableBranchesToShow = branches
+    return [this.currentBranch as string, ...this.selectedBranches]
   }
 
   public getAvailableBranchesToShow() {
@@ -588,7 +589,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   private persistBranchesToShow() {
     return this.persist(
       PersistenceKey.EXPERIMENTS_BRANCHES,
-      this.branchesToShow
+      this.selectedBranches
     )
   }
 

extension/src/experiments/webview/messages.ts

@@ -225,7 +225,7 @@ export class WebviewMessages {
       undefined,
       undefined
     )
-    this.experiments.setBranchesToShow(selectedBranches)
+    this.experiments.setSelectedBranches(selectedBranches)
     await this.update()
   }
 

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

@@ -105,9 +105,7 @@ suite('Experiments Data Test Suite', () => {
             getNbOfCommitsToShow: () => ({
               main: DEFAULT_NUM_OF_COMMITS_TO_SHOW
             }),
-            pruneBranchesToShow: stub(),
-            setAvailableBranchesToShow: stub(),
-            setBranchesToShow: stub()
+            setBranches: stub()
           } as unknown as ExperimentsModel
         )
       )
@@ -168,9 +166,7 @@ suite('Experiments Data Test Suite', () => {
             getNbOfCommitsToShow: () => ({
               main: DEFAULT_NUM_OF_COMMITS_TO_SHOW
             }),
-            pruneBranchesToShow: stub(),
-            setAvailableBranchesToShow: stub(),
-            setBranchesToShow: stub()
+            setBranches: stub()
           } as unknown as ExperimentsModel
         )
       )
@@ -202,23 +198,21 @@ suite('Experiments Data Test Suite', () => {
         'secret-branch',
         'old-branch'
       ]
-      const { data, mockPruneBranchesToShow, mockGetBranchesToShow } =
+      const { data, mockSetBranches, mockGetBranchesToShow } =
         buildExperimentsData(disposable)
       mockGetBranchesToShow.returns(branchesToShow)
 
       await data.update()
 
-      expect(mockPruneBranchesToShow).to.be.calledOnce
+      expect(mockSetBranches).to.be.calledOnce
     })
 
-    it('should add the current branch to the exp show output', async () => {
+    it('should get the required commits from the git log output', async () => {
       stub(ExperimentsData.prototype, 'managedUpdate').resolves()
-      const { data, mockExpShow, mockGetBranchesToShow } = buildExperimentsData(
-        disposable,
-        'branch-283498'
-      )
+      const { data, mockExpShow, mockGetBranchesToShow } =
+        buildExperimentsData(disposable)
 
-      mockGetBranchesToShow.returns([])
+      mockGetBranchesToShow.returns(['main'])
 
       await data.update()
 

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

@@ -1451,10 +1451,13 @@ suite('Experiments Test Suite', () => {
         mockUpdateExperimentsData,
         mockSelectBranches
       } = setupExperimentsAndMockCommands()
-      const mockSetBranchesToShow = stub(experimentsModel, 'setBranchesToShow')
+      const mockSetSelectedBranches = stub(
+        experimentsModel,
+        'setSelectedBranches'
+      )
 
       const waitForBranchesToBeSelected = new Promise(resolve =>
-        mockSetBranchesToShow.callsFake(() => resolve(undefined))
+        mockSetSelectedBranches.callsFake(() => resolve(undefined))
       )
 
       const webview = await experiments.showWebview()
@@ -1469,7 +1472,7 @@ suite('Experiments Test Suite', () => {
 
       await waitForBranchesToBeSelected
 
-      expect(mockSetBranchesToShow).to.be.calledOnceWith(['main', 'other'])
+      expect(mockSetSelectedBranches).to.be.calledOnceWith(['main', 'other'])
 
       expect(mockUpdateExperimentsData).to.be.calledOnce
     }).timeout(WEBVIEW_TEST_TIMEOUT)
@@ -1506,7 +1509,7 @@ suite('Experiments Test Suite', () => {
         mockUpdateExperimentsData,
         mockSelectBranches
       } = setupExperimentsAndMockCommands()
-      const mockSetBranchesToShow = stub(experimentsModel, 'setBranchesToShow')
+      const mockSetBranches = stub(experimentsModel, 'setBranches')
       mockSelectBranches.resolves(undefined)
 
       const webview = await experiments.showWebview()
@@ -1519,7 +1522,10 @@ suite('Experiments Test Suite', () => {
 
       expect(mockSelectBranches).to.be.calledOnce
 
-      expect(mockSetBranchesToShow).not.to.be.calledOnceWith(['main', 'other'])
+      expect(mockSetBranches).not.to.be.calledOnceWith('main', [
+        'main',
+        'other'
+      ])
 
       expect(mockUpdateExperimentsData).not.to.be.calledOnce
     }).timeout(WEBVIEW_TEST_TIMEOUT)

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

@@ -199,14 +199,12 @@ export const buildExperimentsData = (
   stub(gitReader, 'getNumCommits').resolves(404)
 
   const mockGetBranchesToShow = stub().returns(['main'])
-  const mockPruneBranchesToShow = stub()
+  const mockSetBranches = stub()
   const data = disposer.track(
     new ExperimentsData(dvcDemoPath, internalCommands, {
       getBranchesToShow: mockGetBranchesToShow,
       getNbOfCommitsToShow: () => DEFAULT_NUM_OF_COMMITS_TO_SHOW,
-      pruneBranchesToShow: mockPruneBranchesToShow,
-      setAvailableBranchesToShow: stub(),
-      setBranchesToShow: stub()
+      setBranches: mockSetBranches
     } as unknown as ExperimentsModel)
   )
 
@@ -215,7 +213,7 @@ export const buildExperimentsData = (
     mockCreateFileSystemWatcher,
     mockExpShow,
     mockGetBranchesToShow,
-    mockPruneBranchesToShow
+    mockSetBranches
   }
 }
 

extension/src/test/suite/util.ts

@@ -183,11 +183,10 @@ export const buildMockExperimentsData = (update = stub()) =>
   ({
     dispose: stub(),
     dvcYamlChanged: stub(),
-    getCurrentBranchesToShow: () => ['current'],
+    getBranchesToShow: () => ['main'],
     onDidChangeDvcYaml: stub(),
     onDidUpdate: stub(),
-    pruneBranchesToShow: stub(),
-    setCurrentBranchesToShow: stub(),
+    setBranches: stub(),
     update
   } as unknown as ExperimentsData)