Move plots webview message sending into webview messages (#2033)

mattseddon · web-flow · commit 253394cf428e · 2022-07-18T20:47:55.000Z
diff --git a/extension/src/plots/index.ts b/extension/src/plots/index.ts
@@ -1,12 +1,6 @@
 import { join } from 'path'
 import { Event, EventEmitter, Memento } from 'vscode'
-import isEmpty from 'lodash.isempty'
-import {
-  ComparisonPlot,
-  ComparisonRevisionData,
-  PlotsData as TPlotsData,
-  Section
-} from './webview/contract'
+import { PlotsData as TPlotsData } from './webview/contract'
 import { WebviewMessages } from './webview/messages'
 import { PlotsData } from './data'
 import { PlotsModel } from './model'
@@ -20,7 +14,7 @@ import { Resource } from '../resourceLocator'
 import { InternalCommands } from '../commands/internal'
 import { definedAndNonEmpty } from '../util/array'
 import { ExperimentsOutput, TEMP_PLOTS_DIR } from '../cli/reader'
-import { getModifiedTime, removeDir } from '../fileSystem'
+import { removeDir } from '../fileSystem'
 import { Toast } from '../vscode/toast'
 import { pickPaths } from '../path/selection/quickPick'
 
@@ -41,6 +35,8 @@ export class Plots extends BaseRepository<TPlotsData> {
   private readonly data: PlotsData
   private readonly workspaceState: Memento
 
+  private webviewMessages?: WebviewMessages
+
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
@@ -82,7 +78,11 @@ export class Plots extends BaseRepository<TPlotsData> {
       new PathsModel(this.dvcRoot, this.workspaceState)
     )
 
-    this.handleMessageFromWebview(this.paths, this.plots, this.experiments)
+    this.webviewMessages = this.createWebviewMessageHandler(
+      this.paths,
+      this.plots,
+      this.experiments
+    )
 
     this.data.setModel(this.plots)
 
@@ -144,98 +144,21 @@ export class Plots extends BaseRepository<TPlotsData> {
     this.fetchMissingOrSendPlots()
   }
 
-  private sendCheckpointPlotsData() {
-    this.webview?.show({
-      checkpoint: this.getCheckpointPlots()
-    })
-  }
-
-  private getCheckpointPlots() {
-    return this.plots?.getCheckpointPlots() || null
-  }
-
   private async fetchMissingOrSendPlots() {
     await this.isReady()
 
     if (
       this.paths?.hasPaths() &&
       definedAndNonEmpty(this.plots?.getUnfetchedRevisions())
     ) {
-      this.sendCheckpointPlotsData()
+      this.webviewMessages?.sendCheckpointPlotsMessage()
       return this.data.managedUpdate()
     }
 
-    return this.sendPlots()
+    return this.webviewMessages?.sendWebviewMessage()
   }
 
-  private sendPlots() {
-    this.webview?.show({
-      checkpoint: this.getCheckpointPlots(),
-      comparison: this.getComparisonPlots(),
-      hasPlots: !!this.paths?.hasPaths(),
-      hasSelectedPlots: definedAndNonEmpty(this.paths?.getSelected()),
-      sectionCollapsed: this.plots?.getSectionCollapsed(),
-      selectedRevisions: this.plots?.getSelectedRevisionDetails(),
-      template: this.getTemplatePlots()
-    })
-  }
-
-  private getTemplatePlots() {
-    const paths = this.paths?.getTemplateOrder()
-    const plots = this.plots?.getTemplatePlots(paths)
-
-    if (!this.plots || !plots || isEmpty(plots)) {
-      return null
-    }
-
-    return {
-      plots,
-      size: this.plots.getPlotSize(Section.TEMPLATE_PLOTS)
-    }
-  }
-
-  private getComparisonPlots() {
-    const paths = this.paths?.getComparisonPaths()
-    const comparison = this.plots?.getComparisonPlots(paths)
-    if (!this.plots || !comparison || isEmpty(comparison)) {
-      return null
-    }
-
-    return {
-      plots: comparison.map(({ path, revisions }) => {
-        return { path, revisions: this.getRevisionsWithCorrectUrls(revisions) }
-      }),
-      size: this.plots.getPlotSize(Section.COMPARISON_TABLE)
-    }
-  }
-
-  private getRevisionsWithCorrectUrls(revisions: ComparisonRevisionData) {
-    const acc: ComparisonRevisionData = {}
-
-    for (const [revision, plot] of Object.entries(revisions)) {
-      const updatedPlot = this.addCorrectUrl(plot)
-      if (!updatedPlot) {
-        continue
-      }
-      acc[revision] = updatedPlot
-    }
-    return acc
-  }
-
-  private addCorrectUrl(plot: ComparisonPlot) {
-    if (this.webview) {
-      return {
-        ...plot,
-        url: plot.url
-          ? `${this.webview.getWebviewUri(plot.url)}?${getModifiedTime(
-              plot.url
-            )}`
-          : undefined
-      }
-    }
-  }
-
-  private handleMessageFromWebview(
+  private createWebviewMessageHandler(
     paths: PathsModel,
     plots: PlotsModel,
     experiments: Experiments
@@ -244,10 +167,7 @@ export class Plots extends BaseRepository<TPlotsData> {
       paths,
       plots,
       experiments,
-      () => this.sendSectionCollapsed(),
-      () => this.sendTemplatePlots(),
-      () => this.sendComparisonPlots(),
-      () => this.sendCheckpointPlotsData(),
+      () => this.getWebview(),
       () => this.selectPlots(),
       () => this.data.update()
     )
@@ -256,25 +176,7 @@ export class Plots extends BaseRepository<TPlotsData> {
         webviewMessages.handleMessageFromWebview(message)
       )
     )
-  }
-
-  private sendSectionCollapsed() {
-    this.webview?.show({
-      sectionCollapsed: this.plots?.getSectionCollapsed()
-    })
-  }
-
-  private sendComparisonPlots() {
-    this.webview?.show({
-      comparison: this.getComparisonPlots(),
-      selectedRevisions: this.plots?.getSelectedRevisionDetails()
-    })
-  }
-
-  private sendTemplatePlots() {
-    this.webview?.show({
-      template: this.getTemplatePlots()
-    })
+    return webviewMessages
   }
 
   private waitForInitialData(experiments: Experiments) {
diff --git a/extension/src/plots/webview/messages.ts b/extension/src/plots/webview/messages.ts
@@ -1,4 +1,12 @@
-import { PlotSize, Section, SectionCollapsed } from './contract'
+import isEmpty from 'lodash.isempty'
+import {
+  ComparisonPlot,
+  ComparisonRevisionData,
+  PlotsData as TPlotsData,
+  PlotSize,
+  Section,
+  SectionCollapsed
+} from './contract'
 import { Logger } from '../../common/logger'
 import { Experiments } from '../../experiments'
 import { sendTelemetryEvent } from '../../telemetry'
@@ -11,41 +19,53 @@ import {
 } from '../../webview/contract'
 import { PlotsModel } from '../model'
 import { PathsModel } from '../paths/model'
+import { BaseWebview } from '../../webview'
+import { getModifiedTime } from '../../fileSystem'
+import { definedAndNonEmpty } from '../../util/array'
 
 export class WebviewMessages {
   private readonly paths: PathsModel
   private readonly plots: PlotsModel
   private readonly experiments: Experiments
 
-  private readonly sendSectionCollapsed: () => void
-  private readonly sendTemplatePlots: () => void
-  private readonly sendComparisonPlots: () => void
-  private readonly sendCheckpointPlotsData: () => void
+  private readonly getWebview: () => BaseWebview<TPlotsData> | undefined
   private readonly selectPlots: () => void
   private readonly updateData: () => void
 
   constructor(
     paths: PathsModel,
     plots: PlotsModel,
     experiments: Experiments,
-    sendSectionCollapsed: () => void,
-    sendTemplatePlots: () => void,
-    sendComparisonPlots: () => void,
-    sendCheckpointPlotsData: () => void,
+    getWebview: () => BaseWebview<TPlotsData> | undefined,
     selectPlots: () => void,
     updateData: () => void
   ) {
     this.paths = paths
     this.plots = plots
     this.experiments = experiments
-    this.sendSectionCollapsed = sendSectionCollapsed
-    this.sendTemplatePlots = sendTemplatePlots
-    this.sendComparisonPlots = sendComparisonPlots
-    this.sendCheckpointPlotsData = sendCheckpointPlotsData
+    this.getWebview = getWebview
     this.selectPlots = selectPlots
     this.updateData = updateData
   }
 
+  public sendWebviewMessage() {
+    this.getWebview()?.show({
+      checkpoint: this.getCheckpointPlots(),
+      comparison: this.getComparisonPlots(),
+      hasPlots: !!this.paths?.hasPaths(),
+      hasSelectedPlots: definedAndNonEmpty(this.paths.getSelected()),
+      sectionCollapsed: this.plots.getSectionCollapsed(),
+      selectedRevisions: this.plots.getSelectedRevisionDetails(),
+      template: this.getTemplatePlots()
+    })
+  }
+
+  public sendCheckpointPlotsMessage() {
+    this.getWebview()?.show({
+      checkpoint: this.getCheckpointPlots()
+    })
+  }
+
   public handleMessageFromWebview(message: MessageFromWebview) {
     switch (message.type) {
       case MessageFromWebviewType.TOGGLE_METRIC:
@@ -176,7 +196,84 @@ export class WebviewMessages {
       | typeof EventName.VIEWS_REORDER_PLOTS_METRICS
       | typeof EventName.VIEWS_PLOTS_METRICS_SELECTED
   ) {
-    this.sendCheckpointPlotsData()
+    this.sendCheckpointPlotsMessage()
     sendTelemetryEvent(event, undefined, undefined)
   }
+
+  private sendSectionCollapsed() {
+    this.getWebview()?.show({
+      sectionCollapsed: this.plots?.getSectionCollapsed()
+    })
+  }
+
+  private sendComparisonPlots() {
+    this.getWebview()?.show({
+      comparison: this.getComparisonPlots(),
+      selectedRevisions: this.plots?.getSelectedRevisionDetails()
+    })
+  }
+
+  private sendTemplatePlots() {
+    this.getWebview()?.show({
+      template: this.getTemplatePlots()
+    })
+  }
+
+  private getTemplatePlots() {
+    const paths = this.paths?.getTemplateOrder()
+    const plots = this.plots?.getTemplatePlots(paths)
+
+    if (!this.plots || !plots || isEmpty(plots)) {
+      return null
+    }
+
+    return {
+      plots,
+      size: this.plots.getPlotSize(Section.TEMPLATE_PLOTS)
+    }
+  }
+
+  private getComparisonPlots() {
+    const paths = this.paths.getComparisonPaths()
+    const comparison = this.plots.getComparisonPlots(paths)
+    if (!this.plots || !comparison || isEmpty(comparison)) {
+      return null
+    }
+
+    return {
+      plots: comparison.map(({ path, revisions }) => {
+        return { path, revisions: this.getRevisionsWithCorrectUrls(revisions) }
+      }),
+      size: this.plots.getPlotSize(Section.COMPARISON_TABLE)
+    }
+  }
+
+  private getRevisionsWithCorrectUrls(revisions: ComparisonRevisionData) {
+    const acc: ComparisonRevisionData = {}
+
+    for (const [revision, plot] of Object.entries(revisions)) {
+      const updatedPlot = this.addCorrectUrl(plot)
+      if (!updatedPlot) {
+        continue
+      }
+      acc[revision] = updatedPlot
+    }
+    return acc
+  }
+
+  private addCorrectUrl(plot: ComparisonPlot) {
+    const webview = this.getWebview()
+    if (webview) {
+      return {
+        ...plot,
+        url: plot.url
+          ? `${webview.getWebviewUri(plot.url)}?${getModifiedTime(plot.url)}`
+          : undefined
+      }
+    }
+  }
+
+  private getCheckpointPlots() {
+    return this.plots?.getCheckpointPlots() || null
+  }
 }
diff --git a/extension/src/test/suite/plots/index.test.ts b/extension/src/test/suite/plots/index.test.ts
@@ -167,12 +167,11 @@ suite('Plots Test Suite', () => {
 
     it('should re-fetch data when moving between branches', async () => {
       const mockNow = getMockNow()
-      const { data, experiments, mockPlotsDiff, plots, plotsModel } =
+      const { data, experiments, mockPlotsDiff, plotsModel, webviewMessages } =
         await buildPlots(disposable, plotsDiffFixture)
       mockPlotsDiff.resetHistory()
 
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      const mockSendPlots = stub(plots as any, 'sendPlots')
+      const mockSendPlots = stub(webviewMessages, 'sendWebviewMessage')
 
       mockPlotsDiff
         .onFirstCall()
diff --git a/extension/src/test/suite/plots/util.ts b/extension/src/test/suite/plots/util.ts
@@ -18,6 +18,7 @@ import { MOCK_IMAGE_MTIME } from '../../fixtures/plotsDiff'
 import { PathsModel } from '../../../plots/paths/model'
 import { Color } from '../../../experiments/model/status/colors'
 import { BaseWorkspaceWebviews } from '../../../webview/workspace'
+import { WebviewMessages } from '../../../plots/webview/messages'
 
 export const buildPlots = async (
   disposer: Disposer,
@@ -78,6 +79,9 @@ export const buildPlots = async (
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   const pathsModel: PathsModel = (plots as any).paths
 
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const webviewMessages: WebviewMessages = (plots as any).webviewMessages
+
   return {
     data,
     experiments,
@@ -87,7 +91,8 @@ export const buildPlots = async (
     mockRemoveDir,
     pathsModel,
     plots,
-    plotsModel
+    plotsModel,
+    webviewMessages
   }
 }
 
