Add manual refresh button to missing plots (#1754)

* add basic refresh button with action to comparison table

* add ability to refresh a specific revision

* add unit test for button click action

* add box to empty plot

* update toast text

* align images

Author: mattseddon

extension/src/plots/index.ts

@@ -30,6 +30,7 @@ import { ExperimentsOutput, TEMP_PLOTS_DIR } from '../cli/reader'
 import { getModifiedTime, removeDir } from '../fileSystem'
 import { sendTelemetryEvent } from '../telemetry'
 import { EventName } from '../telemetry/constants'
+import { Toast } from '../vscode/toast'
 
 export type PlotsWebview = BaseWebview<TPlotsData>
 
@@ -228,9 +229,11 @@ export class Plots extends BaseRepository<TPlotsData> {
     if (this.webview) {
       return {
         ...plot,
-        url: `${this.webview.getWebviewUri(plot.url)}?${getModifiedTime(
-          plot.url
-        )}`
+        url: plot.url
+          ? `${this.webview.getWebviewUri(plot.url)}?${getModifiedTime(
+              plot.url
+            )}`
+          : undefined
       }
     }
   }
@@ -263,6 +266,8 @@ export class Plots extends BaseRepository<TPlotsData> {
             return this.selectPlotsFromWebview()
           case MessageFromWebviewType.SELECT_EXPERIMENTS:
             return this.selectExperimentsFromWebview()
+          case MessageFromWebviewType.REFRESH_REVISION:
+            return this.attemptToRefreshData(message.payload)
           default:
             Logger.error(`Unexpected message: ${JSON.stringify(message)}`)
         }
@@ -348,6 +353,17 @@ export class Plots extends BaseRepository<TPlotsData> {
     )
   }
 
+  private attemptToRefreshData(revision: string) {
+    Toast.infoWithOptions(`Attempting to refresh ${revision} plots data.`)
+    this.plots?.setupManualRefresh(revision)
+    this.data.managedUpdate()
+    sendTelemetryEvent(
+      EventName.VIEWS_PLOTS_MANUAL_REFRESH,
+      undefined,
+      undefined
+    )
+  }
+
   private sendCheckpointPlotsAndEvent(
     event:
       | typeof EventName.VIEWS_REORDER_PLOTS_METRICS

extension/src/plots/model/index.ts

@@ -157,6 +157,10 @@ export class PlotsModel extends ModelWithPersistence {
     }
   }
 
+  public setupManualRefresh(id: string) {
+    this.deleteRevisionData(id)
+  }
+
   public getUnfetchedRevisions() {
     return this.getSelectedRevisions().filter(
       revision => !this.fetchedRevs.has(revision)
@@ -315,14 +319,18 @@ export class PlotsModel extends ModelWithPersistence {
   private removeStaleBranches() {
     for (const { id, sha } of this.experiments.getBranchRevisions()) {
       if (sha && this.branchRevisions[id] !== sha) {
-        delete this.revisionData[id]
-        delete this.comparisonData[id]
-        this.fetchedRevs.delete(id)
+        this.deleteRevisionData(id)
         this.branchRevisions[id] = sha
       }
     }
   }
 
+  private deleteRevisionData(id: string) {
+    delete this.revisionData[id]
+    delete this.comparisonData[id]
+    this.fetchedRevs.delete(id)
+  }
+
   private getSelectedRevisions() {
     return this.experiments.getSelectedRevisions().map(({ label }) => label)
   }
@@ -369,11 +377,9 @@ export class PlotsModel extends ModelWithPersistence {
 
     for (const revision of selectedRevisions) {
       const image = this.comparisonData?.[revision]?.[path]
-      if (image) {
-        pathRevisions.revisions[revision] = {
-          revision,
-          url: image.url
-        }
+      pathRevisions.revisions[revision] = {
+        revision,
+        url: image?.url
       }
     }
     acc.push(pathRevisions)

extension/src/plots/webview/contract.ts

@@ -126,7 +126,7 @@ export interface TemplatePlotsData {
 }
 
 export type ComparisonPlot = {
-  url: string
+  url: string | undefined
   revision: string
 }
 

extension/src/telemetry/constants.ts

@@ -46,6 +46,7 @@ export const EventName = Object.assign(
     VIEWS_PLOTS_CLOSED: 'views.plots.closed',
     VIEWS_PLOTS_CREATED: 'views.plots.created',
     VIEWS_PLOTS_FOCUS_CHANGED: 'views.plots.focusChanged',
+    VIEWS_PLOTS_MANUAL_REFRESH: 'views.plots.manualRefresh',
     VIEWS_PLOTS_METRICS_SELECTED: 'views.plots.metricsSelected',
     VIEWS_PLOTS_RENAME_SECTION: 'views.plots.sectionRenamed',
     VIEWS_PLOTS_REVISIONS_REORDERED: 'views.plots.revisionsReordered',
@@ -190,6 +191,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_PLOTS_CLOSED]: undefined
   [EventName.VIEWS_PLOTS_CREATED]: undefined
   [EventName.VIEWS_PLOTS_FOCUS_CHANGED]: WebviewFocusChangedProperties
+  [EventName.VIEWS_PLOTS_MANUAL_REFRESH]: undefined
   [EventName.VIEWS_PLOTS_METRICS_SELECTED]: undefined
   [EventName.VIEWS_PLOTS_RENAME_SECTION]: { section: Section }
   [EventName.VIEWS_PLOTS_REVISIONS_REORDERED]: undefined

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

@@ -519,6 +519,42 @@ suite('Plots Test Suite', () => {
       )
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
+    it('should handle a manual refresh message from the webview', async () => {
+      const { data, plots, plotsModel, mockPlotsDiff } = await buildPlots(
+        disposable,
+        plotsDiffFixture
+      )
+
+      const removeDataSpy = spy(plotsModel, 'setupManualRefresh')
+
+      const webview = await plots.showWebview()
+      mockPlotsDiff.resetHistory()
+
+      const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      const dataUpdateEvent = new Promise(resolve =>
+        data.onDidUpdate(() => resolve(undefined))
+      )
+
+      mockMessageReceived.fire({
+        payload: 'main',
+        type: MessageFromWebviewType.REFRESH_REVISION
+      })
+
+      await dataUpdateEvent
+
+      expect(removeDataSpy).to.be.calledOnce
+      expect(mockSendTelemetryEvent).to.be.calledOnce
+      expect(mockSendTelemetryEvent).to.be.calledWithExactly(
+        EventName.VIEWS_PLOTS_MANUAL_REFRESH,
+        undefined,
+        undefined
+      )
+      expect(mockPlotsDiff).to.be.calledOnce
+      expect(mockPlotsDiff).to.be.calledWithExactly(dvcDemoPath, 'main')
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should be able to make the plots webview visible', async () => {
       const { plots, messageSpy, mockPlotsDiff } = await buildPlots(
         disposable,

extension/src/webview/contract.ts

@@ -21,6 +21,7 @@ export enum MessageFromWebviewType {
   REORDER_PLOTS_COMPARISON = 'reorder-plots-comparison',
   REORDER_PLOTS_METRICS = 'reorder-plots-metrics',
   REORDER_PLOTS_TEMPLATES = 'reorder-plots-templates',
+  REFRESH_REVISION = 'refresh-revision',
   RESIZE_COLUMN = 'resize-column',
   RESIZE_PLOTS = 'resize-plots',
   SORT_COLUMN = 'sort-column',
@@ -124,6 +125,7 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.INITIALIZED }
   | { type: MessageFromWebviewType.SELECT_EXPERIMENTS }
   | { type: MessageFromWebviewType.SELECT_PLOTS }
+  | { type: MessageFromWebviewType.REFRESH_REVISION; payload: string }
 
 export type MessageToWebview<T extends WebviewData> = {
   type: MessageToWebviewType.SET_DATA

webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx

@@ -223,6 +223,39 @@ describe('ComparisonTable', () => {
     expect(pinnedColumn.textContent).toBe(pinnedRevision)
   })
 
+  it('should display a refresh button for each revision that has a missing image', () => {
+    const revisionWithNoData = 'missing-data'
+
+    renderTable({
+      ...basicProps,
+      plots: basicProps.plots.map(({ path, revisions }) => ({
+        path,
+        revisions: {
+          ...revisions,
+          [revisionWithNoData]: { revision: revisionWithNoData, url: undefined }
+        }
+      })),
+      revisions: [
+        ...basicProps.revisions,
+        { displayColor: '#f56565', revision: revisionWithNoData }
+      ]
+    })
+
+    const refreshButtons = screen.getAllByText('Refresh')
+
+    expect(refreshButtons).toHaveLength(basicProps.plots.length)
+
+    for (const button of refreshButtons) {
+      fireEvent.click(button)
+      expect(mockPostMessage).toBeCalledTimes(1)
+      expect(mockPostMessage).toBeCalledWith({
+        payload: revisionWithNoData,
+        type: MessageFromWebviewType.REFRESH_REVISION
+      })
+      mockPostMessage.mockReset()
+    }
+  })
+
   describe('Columns drag and drop', () => {
     const pinSecondColumn = () => {
       const secondColumn = getPin(screen.getByText(revisions[1]))

webview/src/plots/components/comparisonTable/ComparisonTableRow.test.tsx

@@ -12,6 +12,8 @@ import {
 } from './ComparisonTableRow'
 import styles from '../styles.module.scss'
 
+jest.mock('../../../shared/api')
+
 describe('ComparisonTableRow', () => {
   afterEach(() => {
     cleanup()

webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx

@@ -1,8 +1,11 @@
 import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
+import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 import React, { useState } from 'react'
 import cx from 'classnames'
 import styles from './styles.module.scss'
 import { AllIcons, Icon } from '../../../shared/components/Icon'
+import { RefreshButton } from '../../../shared/components/button/RefreshButton'
+import { sendMessage } from '../../../shared/vscode'
 
 export interface ComparisonTableRowProps {
   path: string
@@ -35,23 +38,33 @@ export const ComparisonTableRow: React.FC<ComparisonTableRowProps> = ({
         {nbColumns > 1 && <td colSpan={nbColumns - 1}></td>}
       </tr>
       <tr>
-        {plots.map((plot: ComparisonPlot | undefined) => {
-          const isPinned = pinnedColumn === plot?.revision
+        {plots.map((plot: ComparisonPlot) => {
+          const isPinned = pinnedColumn === plot.revision
           const missing = !plot?.url
 
           return (
             <td
-              key={path + plot?.revision || 'missing'}
+              key={path + plot.revision}
               className={cx({
                 [styles.pinnedColumnCell]: isPinned,
-                [styles.missing]: missing
+                [styles.missing]: isShown && missing
               })}
             >
               <div
                 className={cx(styles.cell, { [styles.cellHidden]: !isShown })}
               >
                 {missing ? (
-                  <p>No plot to display.</p>
+                  <div>
+                    <p>No Plot to Display.</p>
+                    <RefreshButton
+                      onClick={() =>
+                        sendMessage({
+                          payload: plot.revision,
+                          type: MessageFromWebviewType.REFRESH_REVISION
+                        })
+                      }
+                    />
+                  </div>
                 ) : (
                   <img
                     src={plot.url}

webview/src/plots/components/comparisonTable/styles.module.scss

@@ -83,6 +83,9 @@ $gap: 4px;
 
 .missing {
   background-color: $bg-color;
+  border-style: solid;
+  border-width: thin;
+  border-color: $watermark-color;
 }
 
 .rowToggler {
@@ -148,6 +151,7 @@ $gap: 4px;
 .cell img {
   width: 100%;
   height: auto;
+  vertical-align: middle;
 }
 
 .experimentName {