Cleanup plots workspace state (#3518)

Author: julieg18

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

@@ -1,5 +1,6 @@
 import { PlotsModel } from '.'
 import {
+  CustomPlotType,
   DEFAULT_NB_ITEMS_PER_ROW,
   DEFAULT_SECTION_COLLAPSED,
   DEFAULT_SECTION_NB_ITEMS_PER_ROW_OR_WIDTH,
@@ -22,7 +23,7 @@ const mockedRevisions = [
 describe('plotsModel', () => {
   let model: PlotsModel
   const exampleDvcRoot = 'test'
-  const memento = buildMockMemento({
+  let memento = buildMockMemento({
     [PersistenceKey.PLOTS_CUSTOM_ORDER + exampleDvcRoot]:
       customPlotsOrderFixture,
     [PersistenceKey.PLOT_NB_ITEMS_PER_ROW_OR_WIDTH + exampleDvcRoot]:
@@ -45,6 +46,41 @@ describe('plotsModel', () => {
     jest.clearAllMocks()
   })
 
+  it('should update outdated custom and trends state', () => {
+    memento = buildMockMemento({
+      [PersistenceKey.PLOTS_CUSTOM_ORDER + exampleDvcRoot]: [
+        {
+          metric: 'metrics:summary.json:loss',
+          param: 'params:params.yaml:dropout'
+        }
+      ],
+      [PersistenceKey.PLOT_SELECTED_METRICS + exampleDvcRoot]: ['string'],
+      [PersistenceKey.PLOT_METRIC_ORDER + exampleDvcRoot]: ['string']
+    })
+    model = new PlotsModel(
+      exampleDvcRoot,
+      {
+        getFirstThreeColumnOrder: mockedGetFirstThreeColumnOrder,
+        getSelectedRevisions: mockedGetSelectedRevisions,
+        isReady: () => Promise.resolve(undefined)
+      } as unknown as Experiments,
+      memento
+    )
+    expect(model.getCustomPlotsOrder()).toStrictEqual([
+      {
+        metric: 'summary.json:loss',
+        param: 'params.yaml:dropout',
+        type: CustomPlotType.METRIC_VS_PARAM
+      }
+    ])
+    expect(
+      memento.get(PersistenceKey.PLOT_SELECTED_METRICS + exampleDvcRoot)
+    ).toStrictEqual(undefined)
+    expect(
+      memento.get(PersistenceKey.PLOT_METRIC_ORDER + exampleDvcRoot)
+    ).toStrictEqual(undefined)
+  })
+
   it('should change the plotSize when calling setPlotSize', () => {
     expect(
       model.getNbItemsPerRowOrWidth(PlotsSection.CUSTOM_PLOTS)

extension/src/plots/model/index.ts

@@ -91,6 +91,9 @@ export class PlotsModel extends ModelWithPersistence {
     )
     this.comparisonOrder = this.revive(PersistenceKey.PLOT_COMPARISON_ORDER, [])
     this.customPlotsOrder = this.revive(PersistenceKey.PLOTS_CUSTOM_ORDER, [])
+
+    this.cleanupOutdatedCustomPlotsState()
+    this.cleanupOutdatedTrendsState()
   }
 
   public transformAndSetExperiments() {
@@ -181,9 +184,7 @@ export class PlotsModel extends ModelWithPersistence {
   }
 
   public getCustomPlotsOrder() {
-    return this.customPlotsOrder.map(value =>
-      cleanupOldOrderValue(value, FILE_SEPARATOR)
-    )
+    return this.customPlotsOrder
   }
 
   public updateCustomPlotsOrder(plotsOrder: CustomPlotsOrderValue[]) {
@@ -401,6 +402,24 @@ export class PlotsModel extends ModelWithPersistence {
     return mapping
   }
 
+  private cleanupOutdatedCustomPlotsState() {
+    const order = this.getCustomPlotsOrder()
+    const workspaceHoldsUpToDateState =
+      order.length === 0 || order[0].type !== undefined
+    if (workspaceHoldsUpToDateState) {
+      return
+    }
+    const newOrder = order.map(value =>
+      cleanupOldOrderValue(value, FILE_SEPARATOR)
+    )
+    this.setCustomPlotsOrder(newOrder)
+  }
+
+  private cleanupOutdatedTrendsState() {
+    this.persist(PersistenceKey.PLOT_METRIC_ORDER, undefined)
+    this.persist(PersistenceKey.PLOT_SELECTED_METRICS, undefined)
+  }
+
   private removeStaleData() {
     return Promise.all([this.removeStaleCommits(), this.removeStaleRevisions()])
   }