Save smooth plot values across sessions (#4220)

Author: julieg18

extension/src/persistence/constants.ts

@@ -17,6 +17,7 @@ export enum PersistenceKey {
   PLOTS_CUSTOM_ORDER = 'plotCustomOrder:',
   PLOT_SECTION_COLLAPSED = 'plotSectionCollapsed:',
   PLOT_SELECTED_METRICS = 'plotSelectedMetrics:',
+  PLOTS_SMOOTH_PLOT_VALUES = 'plotSmoothPlotValues:',
   PLOT_TEMPLATE_ORDER = 'plotTemplateOrder:'
 }
 

extension/src/plots/model/index.ts

@@ -32,7 +32,8 @@ import {
   CustomPlotsData,
   DEFAULT_HEIGHT,
   DEFAULT_NB_ITEMS_PER_ROW,
-  PlotHeight
+  PlotHeight,
+  SmoothPlotValues
 } from '../webview/contract'
 import {
   EXPERIMENT_WORKSPACE_ID,
@@ -74,6 +75,7 @@ export class PlotsModel extends ModelWithPersistence {
 
   private comparisonData: ComparisonData = {}
   private comparisonOrder: string[]
+  private smoothPlotValues: SmoothPlotValues = {}
 
   private revisionData: RevisionData = {}
   private templates: TemplateAccumulator = {}
@@ -102,6 +104,10 @@ export class PlotsModel extends ModelWithPersistence {
     )
     this.comparisonOrder = this.revive(PersistenceKey.PLOT_COMPARISON_ORDER, [])
     this.customPlotsOrder = this.revive(PersistenceKey.PLOTS_CUSTOM_ORDER, [])
+    this.smoothPlotValues = this.revive(
+      PersistenceKey.PLOTS_SMOOTH_PLOT_VALUES,
+      {}
+    )
 
     this.cleanupOutdatedCustomPlotsState()
     this.cleanupOutdatedTrendsState()
@@ -312,6 +318,15 @@ export class PlotsModel extends ModelWithPersistence {
     )
   }
 
+  public setSmoothPlotValues(id: string, value: number) {
+    this.smoothPlotValues[id] = value
+    this.persist(PersistenceKey.PLOTS_SMOOTH_PLOT_VALUES, this.smoothPlotValues)
+  }
+
+  public getSmoothPlotValues() {
+    return this.smoothPlotValues
+  }
+
   public getNbItemsPerRowOrWidth(section: PlotsSection) {
     if (this.nbItemsPerRowOrWidth[section]) {
       return this.nbItemsPerRowOrWidth[section]

extension/src/plots/webview/contract.ts

@@ -141,10 +141,13 @@ export type TemplatePlotSection = {
   entries: TemplatePlotEntry[]
 }
 
+export type SmoothPlotValues = { [id: string]: number }
+
 export interface TemplatePlotsData {
   plots: TemplatePlotSection[]
   nbItemsPerRow: number
   height: PlotHeight
+  smoothPlotValues: SmoothPlotValues
 }
 
 export type ComparisonPlot = {

extension/src/plots/webview/messages.ts

@@ -118,6 +118,11 @@ export class WebviewMessages {
         )
       case MessageFromWebviewType.TOGGLE_EXPERIMENT:
         return this.setExperimentStatus(message.payload)
+      case MessageFromWebviewType.SET_SMOOTH_PLOT_VALUE:
+        return this.setSmoothPlotValues(
+          message.payload.id,
+          message.payload.value
+        )
       case MessageFromWebviewType.ZOOM_PLOT:
         if (message.payload) {
           const imagePath = this.revertCorrectUrl(message.payload)
@@ -200,6 +205,16 @@ export class WebviewMessages {
     )
   }
 
+  private setSmoothPlotValues(id: string, value: number) {
+    this.plots.setSmoothPlotValues(id, value)
+    this.sendTemplatePlots()
+    sendTelemetryEvent(
+      EventName.VIEWS_PLOTS_SET_SMOOTH_PLOT_VALUE,
+      undefined,
+      undefined
+    )
+  }
+
   private setCustomPlotsOrder(plotIds: string[]) {
     const customPlotsOrderWithId = this.plots
       .getCustomPlotsOrder()
@@ -287,7 +302,8 @@ export class WebviewMessages {
       nbItemsPerRow: this.plots.getNbItemsPerRowOrWidth(
         PlotsSection.TEMPLATE_PLOTS
       ),
-      plots
+      plots,
+      smoothPlotValues: this.plots.getSmoothPlotValues()
     }
   }
 

extension/src/telemetry/constants.ts

@@ -79,6 +79,7 @@ export const EventName = Object.assign(
     VIEWS_PLOTS_SECTION_TOGGLE: 'views.plots.toggleSection',
     VIEWS_PLOTS_SELECT_EXPERIMENTS: 'view.plots.selectExperiments',
     VIEWS_PLOTS_SELECT_PLOTS: 'view.plots.selectPlots',
+    VIEWS_PLOTS_SET_SMOOTH_PLOT_VALUE: 'view.plots.setSmoothPlotValues',
     VIEWS_PLOTS_ZOOM_PLOT: 'views.plots.zoomPlot',
     VIEWS_REORDER_PLOTS_CUSTOM: 'views.plots.customReordered',
     VIEWS_REORDER_PLOTS_TEMPLATES: 'views.plots.templatesReordered',
@@ -274,6 +275,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_PLOTS_ZOOM_PLOT]: { isImage: boolean }
   [EventName.VIEWS_REORDER_PLOTS_CUSTOM]: undefined
   [EventName.VIEWS_REORDER_PLOTS_TEMPLATES]: undefined
+  [EventName.VIEWS_PLOTS_SET_SMOOTH_PLOT_VALUE]: undefined
 
   [EventName.VIEWS_PLOTS_PATH_TREE_OPENED]: DvcRootCount
 

extension/src/test/fixtures/plotsDiff/index.ts

@@ -733,13 +733,15 @@ export const getMinimalWebviewMessage = () => ({
   plots: extendedSpecs(basicVega),
   nbItemsPerRow: DEFAULT_NB_ITEMS_PER_ROW,
   height: DEFAULT_PLOT_HEIGHT,
-  revisions: getRevisions()
+  revisions: getRevisions(),
+  smoothPlotValues: {}
 })
 
 export const getTemplateWebviewMessage = (): TemplatePlotsData => ({
   plots: extendedSpecs({ ...basicVega, ...require('./vega').default }),
   nbItemsPerRow: DEFAULT_NB_ITEMS_PER_ROW,
-  height: DEFAULT_PLOT_HEIGHT
+  height: DEFAULT_PLOT_HEIGHT,
+  smoothPlotValues: {}
 })
 
 export const getManyTemplatePlotsWebviewMessage = (
@@ -749,7 +751,8 @@ export const getManyTemplatePlotsWebviewMessage = (
     ...multipleVega(length)
   }),
   nbItemsPerRow: DEFAULT_NB_ITEMS_PER_ROW,
-  height: DEFAULT_PLOT_HEIGHT
+  height: DEFAULT_PLOT_HEIGHT,
+  smoothPlotValues: {}
 })
 
 export const MOCK_IMAGE_MTIME = 946684800000

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

@@ -1038,6 +1038,39 @@ suite('Plots Test Suite', () => {
       )
     })
 
+    it('should handle an update smooth plot values message from the webview', async () => {
+      const { plots, plotsModel } = await buildPlots({
+        disposer: disposable,
+        plotsDiff: plotsDiffFixture
+      })
+      const templatePlot = templatePlotsFixture.plots[0].entries[0]
+
+      const webview = await plots.showWebview()
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+      const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockSetSmoothPlotValues = stub(plotsModel, 'setSmoothPlotValues')
+
+      mockMessageReceived.fire({
+        payload: {
+          id: templatePlot.id,
+          value: 0.5
+        },
+        type: MessageFromWebviewType.SET_SMOOTH_PLOT_VALUE
+      })
+
+      expect(mockSendTelemetryEvent).to.be.called
+      expect(mockSendTelemetryEvent).to.be.calledWithExactly(
+        EventName.VIEWS_PLOTS_SET_SMOOTH_PLOT_VALUE,
+        undefined,
+        undefined
+      )
+      expect(mockSetSmoothPlotValues).to.be.called
+      expect(mockSetSmoothPlotValues).to.be.calledWithExactly(
+        templatePlot.id,
+        0.5
+      )
+    })
+
     it('should handle the CLI throwing an error', async () => {
       const { data, errorsModel, mockPlotsDiff, plots, plotsModel } =
         await buildPlots({ disposer: disposable, plotsDiff: plotsDiffFixture })

extension/src/webview/contract.ts

@@ -40,6 +40,7 @@ export enum MessageFromWebviewType {
   RESIZE_COLUMN = 'resize-column',
   RESIZE_PLOTS = 'resize-plots',
   SAVE_STUDIO_TOKEN = 'save-studio-token',
+  SET_SMOOTH_PLOT_VALUE = 'update-smooth-plot-value',
   SHOW_EXPERIMENT_LOGS = 'show-experiment-logs',
   SHOW_WALKTHROUGH = 'show-walkthrough',
   STOP_EXPERIMENTS = 'stop-experiments',
@@ -232,6 +233,10 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.SHOW_WALKTHROUGH }
   | { type: MessageFromWebviewType.SHOW_SCM_PANEL }
   | { type: MessageFromWebviewType.INSTALL_DVC }
+  | {
+      type: MessageFromWebviewType.SET_SMOOTH_PLOT_VALUE
+      payload: { id: string; value: number }
+    }
   | { type: MessageFromWebviewType.UPGRADE_DVC }
   | { type: MessageFromWebviewType.SETUP_WORKSPACE }
   | { type: MessageFromWebviewType.OPEN_STUDIO }

webview/src/plots/components/App.test.tsx

@@ -2147,5 +2147,123 @@ describe('App', () => {
       fireEvent(panel, clickEvent)
       expect(clickEvent.stopPropagation).toHaveBeenCalledTimes(1)
     })
+
+    describe('Smooth Plots', () => {
+      const waitSetValuePostMessage = (value: number) =>
+        waitFor(
+          () =>
+            expect(mockPostMessage).toHaveBeenCalledWith({
+              payload: { id: smoothId, value },
+              type: MessageFromWebviewType.SET_SMOOTH_PLOT_VALUE
+            }),
+          { timeout: 5000 }
+        )
+      it('should send a message to save the value when a vega panel slider is interacted with', async () => {
+        renderAppWithOptionalData({ template: withVegaPanels })
+
+        const smoothPlot = screen.getByTestId(`plot_${smoothId}`)
+        await waitForVega(smoothPlot)
+
+        // eslint-disable-next-line testing-library/no-node-access
+        const slider = smoothPlot.querySelector(
+          '.vega-bindings input[name="smooth"]'
+        )
+        expect(slider).toBeInTheDocument()
+
+        fireEvent.change(slider as HTMLInputElement, { target: { value: 0.4 } })
+
+        await waitSetValuePostMessage(0.4)
+      })
+
+      it('should send a message to save the value when a zoomed in plot vega panel slider is interacted with', async () => {
+        renderAppWithOptionalData({ template: withVegaPanels })
+
+        const smoothPlot = within(
+          screen.getByTestId(`plot_${smoothId}`)
+        ).getByRole('button')
+        fireEvent.click(smoothPlot)
+
+        const popup = screen.getByTestId('zoomed-in-plot')
+        await waitForVega(popup)
+
+        // eslint-disable-next-line testing-library/no-node-access
+        const slider = popup.querySelector(
+          '.vega-bindings input[name="smooth"]'
+        )
+        expect(slider).toBeInTheDocument()
+
+        fireEvent.change(slider as HTMLInputElement, { target: { value: 0.4 } })
+
+        await waitSetValuePostMessage(0.4)
+      })
+
+      it('should set a vega panel slider value when given a default value', async () => {
+        renderAppWithOptionalData({
+          template: { ...withVegaPanels, smoothPlotValues: { [smoothId]: 0.6 } }
+        })
+
+        const smoothPlot = screen.getByTestId(`plot_${smoothId}`)
+        await waitForVega(smoothPlot)
+
+        // eslint-disable-next-line testing-library/no-node-access
+        const slider = smoothPlot.querySelector(
+          '.vega-bindings input[name="smooth"]'
+        )
+        expect(slider).toBeInTheDocument()
+
+        expect(slider).toHaveValue('0.6')
+      })
+
+      it('should set the zoomed in plot vega panel slider value when given a default value', async () => {
+        renderAppWithOptionalData({
+          template: { ...withVegaPanels, smoothPlotValues: { [smoothId]: 0.6 } }
+        })
+
+        const smoothPlot = within(
+          screen.getByTestId(`plot_${smoothId}`)
+        ).getByRole('button')
+        fireEvent.click(smoothPlot)
+
+        const popup = screen.getByTestId('zoomed-in-plot')
+        await waitForVega(popup)
+
+        // eslint-disable-next-line testing-library/no-node-access
+        const slider = popup.querySelector(
+          '.vega-bindings input[name="smooth"]'
+        )
+        expect(slider).toBeInTheDocument()
+
+        expect(slider).toHaveValue('0.6')
+      })
+
+      it('should update a vega panel slider value when given a new value', async () => {
+        renderAppWithOptionalData({
+          template: { ...withVegaPanels }
+        })
+
+        const smoothPlot = screen.getByTestId(`plot_${smoothId}`)
+
+        await waitForVega(smoothPlot)
+
+        // eslint-disable-next-line testing-library/no-node-access
+        const slider = smoothPlot.querySelector(
+          '.vega-bindings input[name="smooth"]'
+        )
+
+        expect(slider).toBeInTheDocument()
+        expect(slider).toHaveValue('0.2')
+
+        sendSetDataMessage({
+          template: {
+            ...withVegaPanels,
+            smoothPlotValues: { [smoothId]: 0.7 }
+          }
+        })
+
+        await waitFor(() => expect(slider).toHaveValue('0.7'), {
+          timeout: 5000
+        })
+      })
+    })
   })
 })

webview/src/plots/components/Plots.tsx

@@ -55,7 +55,11 @@ const PlotsContent = () => {
         dispatch(setZoomedInPlot(undefined))
       }}
     >
-      <ZoomedInPlot id={zoomedInPlot.id} props={zoomedInPlot.plot} />
+      <ZoomedInPlot
+        isTemplatePlot={zoomedInPlot.isTemplatePlot}
+        id={zoomedInPlot.id}
+        props={zoomedInPlot.plot}
+      />
     </Modal>
   )