Add constant to verify if there is a custom selection of plots (#4989)

* Add constant to check if there previously was a custom selection of plots

* Add tests

* Added more testsM

* Stub other functions

Author: sroy3

extension/src/persistence/constants.ts

@@ -15,6 +15,7 @@ export enum PersistenceKey {
   PLOT_METRIC_ORDER = 'plotMetricOrder:',
   PLOT_NB_ITEMS_PER_ROW_OR_WIDTH = 'plotNbItemsPerRowOrWidth:',
   PLOTS_CUSTOM_ORDER = 'plotCustomOrder:',
+  PLOTS_HAS_CUSTOM_SELECTION = 'plotsCustomCollection:',
   PLOT_SECTION_COLLAPSED = 'plotSectionCollapsed:',
   PLOT_SELECTED_METRICS = 'plotSelectedMetrics:',
   PLOTS_SMOOTH_PLOT_VALUES = 'plotSmoothPlotValues:',

extension/src/plots/index.ts

@@ -91,6 +91,7 @@ export class Plots extends BaseRepository<TPlotsData> {
     const status = this.paths.toggleStatus(path)
     this.paths.setTemplateOrder()
     this.notifyChanged()
+    this.setHasCustomSelection()
     return status
   }
 
@@ -104,6 +105,7 @@ export class Plots extends BaseRepository<TPlotsData> {
 
     this.paths.setSelected(selected)
     this.paths.setTemplateOrder()
+    this.setHasCustomSelection()
     return this.notifyChanged()
   }
 
@@ -170,6 +172,10 @@ export class Plots extends BaseRepository<TPlotsData> {
     return this.sendPlots()
   }
 
+  private setHasCustomSelection() {
+    this.paths.setHasCustomSelection(this.paths.getTerminalNodes().length > 20)
+  }
+
   private notifyChanged() {
     const selectedRevisions = this.plots.getSelectedRevisionIds()
     this.paths.setSelectedRevisions(selectedRevisions)
@@ -244,6 +250,7 @@ export class Plots extends BaseRepository<TPlotsData> {
         collectInitialIdShas()
         this.setupExperimentsListener(experiments)
         void this.initializeData()
+        this.paths.checkIfHasPreviousCustomSelection()
       })
     )
   }

extension/src/plots/paths/model.test.ts

@@ -8,6 +8,8 @@ import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 import { ErrorsModel } from '../errors/model'
 import { REVISIONS } from '../../test/fixtures/plotsDiff'
 import { SpecWithTitles } from '../vega/util'
+import { PersistenceKey } from '../../persistence/constants'
+import { Status } from '../../path/selection/model'
 
 describe('PathsModel', () => {
   const mockDvcRoot = 'test'
@@ -502,4 +504,75 @@ describe('PathsModel', () => {
 
     expect(model.getTerminalNodes()).toStrictEqual([])
   })
+
+  it('should not change hasCustomSelection when checking for if it is already defined', () => {
+    const memento = buildMockMemento({
+      [PersistenceKey.PLOTS_HAS_CUSTOM_SELECTION + mockDvcRoot]: true
+    })
+    const model = new PathsModel(mockDvcRoot, buildMockErrorsModel(), memento)
+    const setHasCustomSelectionSpy = jest.spyOn(model, 'setHasCustomSelection')
+
+    model.checkIfHasPreviousCustomSelection()
+    expect(setHasCustomSelectionSpy).not.toHaveBeenCalled()
+  })
+
+  it('should set hasCustomSelection to false if there are less than 20 plots', () => {
+    const model = new PathsModel(
+      mockDvcRoot,
+      buildMockErrorsModel(),
+      buildMockMemento()
+    )
+    const setHasCustomSelectionSpy = jest.spyOn(model, 'setHasCustomSelection')
+    jest
+      .spyOn(model, 'getTerminalNodeStatuses')
+      .mockImplementation(() => [
+        Status.SELECTED,
+        Status.SELECTED,
+        Status.SELECTED,
+        Status.SELECTED
+      ])
+
+    model.checkIfHasPreviousCustomSelection()
+    expect(setHasCustomSelectionSpy).toHaveBeenCalledWith(false)
+  })
+
+  it('should set hasCustomSelection to false if there are 20 selected plots', () => {
+    const model = new PathsModel(
+      mockDvcRoot,
+      buildMockErrorsModel(),
+      buildMockMemento()
+    )
+    const setHasCustomSelectionSpy = jest.spyOn(model, 'setHasCustomSelection')
+    const statuses = [Status.UNSELECTED]
+    for (let i = 0; i < 20; i++) {
+      statuses.push(Status.SELECTED)
+    }
+    statuses.push(Status.UNSELECTED)
+    jest
+      .spyOn(model, 'getTerminalNodeStatuses')
+      .mockImplementation(() => statuses)
+
+    model.checkIfHasPreviousCustomSelection()
+    expect(setHasCustomSelectionSpy).toHaveBeenCalledWith(false)
+  })
+
+  it('should set hasCustomSelection to true if there are more than 20 plots and the number of selected plots is not 20', () => {
+    const model = new PathsModel(
+      mockDvcRoot,
+      buildMockErrorsModel(),
+      buildMockMemento()
+    )
+    const setHasCustomSelectionSpy = jest.spyOn(model, 'setHasCustomSelection')
+    const statuses = [Status.UNSELECTED]
+    for (let i = 0; i < 19; i++) {
+      statuses.push(Status.SELECTED)
+    }
+    statuses.push(Status.UNSELECTED)
+    jest
+      .spyOn(model, 'getTerminalNodeStatuses')
+      .mockImplementation(() => statuses)
+
+    model.checkIfHasPreviousCustomSelection()
+    expect(setHasCustomSelectionSpy).toHaveBeenCalledWith(true)
+  })
 })

extension/src/plots/paths/model.ts

@@ -27,6 +27,8 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
 
   private selectedRevisions: string[] = []
 
+  private hasCustomSelection: boolean | undefined = undefined
+
   constructor(dvcRoot: string, errors: ErrorsModel, workspaceState: Memento) {
     super(dvcRoot, workspaceState, PersistenceKey.PLOT_PATH_STATUS)
 
@@ -37,6 +39,11 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
       PersistenceKey.PLOT_COMPARISON_PATHS_ORDER,
       []
     )
+
+    this.hasCustomSelection = this.revive(
+      PersistenceKey.PLOTS_HAS_CUSTOM_SELECTION,
+      undefined
+    )
   }
 
   public setSelectedRevisions(selectedRevisions: string[]) {
@@ -144,6 +151,30 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
     return this.data.length > 0
   }
 
+  public setHasCustomSelection(hasCustomSelection: boolean) {
+    this.hasCustomSelection = hasCustomSelection
+    this.persist(
+      PersistenceKey.PLOTS_HAS_CUSTOM_SELECTION,
+      this.hasCustomSelection
+    )
+  }
+
+  public checkIfHasPreviousCustomSelection() {
+    if (this.hasCustomSelection === undefined) {
+      const statuses = this.getTerminalNodeStatuses()
+      const plotsLength = statuses.length
+
+      const hasCustomSelection =
+        plotsLength > 20
+          ? plotsLength -
+              statuses.filter(nodeStatus => nodeStatus !== Status.SELECTED)
+                .length !==
+            20
+          : false
+      this.setHasCustomSelection(hasCustomSelection)
+    }
+  }
+
   private handleCliError() {
     this.data = []
   }

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

@@ -55,10 +55,16 @@ import { BaseWebview } from '../../../webview'
 import * as PlotsCollectUtils from '../../../plots/model/collect'
 import { Operator } from '../../../experiments/model/filterBy'
 import * as External from '../../../vscode/external'
+import { PlotPath } from '../../../plots/paths/collect'
 
 suite('Plots Test Suite', () => {
   const disposable = Disposable.fn()
 
+  const createTerminalNodesArray = (length: number) =>
+    Array.from({ length }, i => ({ path: i })) as unknown as (PlotPath & {
+      selected: boolean
+    })[]
+
   beforeEach(() => {
     restore()
   })
@@ -1304,5 +1310,35 @@ suite('Plots Test Suite', () => {
         'should no long provide decorations to the plots paths tree'
       ).to.deep.equal(new Set([]))
     })
+
+    it('should togglePathStatus to true when calling selectPlots with more than 20 plots', async () => {
+      const { plots, pathsModel } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+      const mockSetCustomSelection = stub(pathsModel, 'setHasCustomSelection')
+      stub(pathsModel, 'setTemplateOrder')
+      stub(pathsModel, 'toggleStatus')
+      stub(pathsModel, 'getTerminalNodes').returns(createTerminalNodesArray(32))
+
+      plots.togglePathStatus(dvcDemoPath)
+
+      expect(mockSetCustomSelection).to.be.calledWith(true)
+    })
+
+    it('should togglePathStatus to false when calling selectPlots with 20 plots or less', async () => {
+      const { plots, pathsModel } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+      const mockSetCustomSelection = stub(pathsModel, 'setHasCustomSelection')
+      stub(pathsModel, 'setTemplateOrder')
+      stub(pathsModel, 'toggleStatus')
+      stub(pathsModel, 'getTerminalNodes').returns(createTerminalNodesArray(20))
+
+      plots.togglePathStatus(dvcDemoPath)
+
+      expect(mockSetCustomSelection).to.be.calledWith(false)
+    })
   })
 })