collect plots webview message async

extension/src/experiments/model/index.ts

@@ -42,7 +42,7 @@ import {
 
 type StarredExperiments = Record<string, boolean | undefined>
 
-export type SelectedExperimentWithColor = Experiment & {
+type SelectedExperimentWithColor = Experiment & {
   displayColor: Color
   selected: true
 }

extension/src/plots/webview/messages.ts

@@ -61,18 +61,12 @@ export class WebviewMessages {
   }
 
   public async sendWebviewMessage() {
-    const selectedRevisions = this.plots.getSelectedRevisionDetails()
-
-    await this.getWebview()?.show({
-      cliError: this.errors.getCliError()?.error || null,
-      comparison: this.getComparisonPlots(),
-      custom: this.getCustomPlots(),
-      hasPlots: !!this.paths.hasPaths(),
-      hasUnselectedPlots: this.paths.getHasUnselectedPlots(),
-      sectionCollapsed: this.plots.getSectionCollapsed(),
-      selectedRevisions,
-      template: this.getTemplatePlots(selectedRevisions)
-    })
+    const webview = this.getWebview()
+    if (!webview) {
+      return
+    }
+    const data = await this.getWebviewData()
+    return webview.show(data)
   }
 
   public handleMessageFromWebview(message: MessageFromWebview) {
@@ -140,6 +134,39 @@ export class WebviewMessages {
     }
   }
 
+  private async getWebviewData(): Promise<TPlotsData> {
+    const selectedRevisions = this.plots.getSelectedRevisionDetails()
+
+    const [
+      cliError,
+      comparison,
+      custom,
+      hasPlots,
+      hasUnselectedPlots,
+      sectionCollapsed,
+      template
+    ] = await Promise.all([
+      this.errors.getCliError()?.error || null,
+      this.getComparisonPlots(),
+      this.getCustomPlots(),
+      !!this.paths.hasPaths(),
+      this.paths.getHasUnselectedPlots(),
+      this.plots.getSectionCollapsed(),
+      this.getTemplatePlots(selectedRevisions)
+    ])
+
+    return {
+      cliError,
+      comparison,
+      custom,
+      hasPlots,
+      hasUnselectedPlots,
+      sectionCollapsed,
+      selectedRevisions,
+      template
+    }
+  }
+
   private setPlotSize(
     section: PlotsSection,
     nbItemsPerRow: number,

extension/src/test/suite/experiments/model/tree.test.ts

@@ -16,14 +16,15 @@ import {
   getFirstArgOfLastCall,
   getMockNow,
   getTimeSafeDisposer,
-  stubPrivatePrototypeMethod
+  stubPrivatePrototypeMethod,
+  waitForSpyCall
 } from '../../util'
 import { dvcDemoPath } from '../../../util'
 import {
   RegisteredCliCommands,
   RegisteredCommands
 } from '../../../../commands/external'
-import { buildPlots } from '../../plots/util'
+import { buildPlots, buildPlotsWebview } from '../../plots/util'
 import { ExperimentsTree } from '../../../../experiments/model/tree'
 import { buildExperiments, stubWorkspaceExperimentsGetters } from '../util'
 import { WEBVIEW_TEST_TIMEOUT } from '../../timeouts'
@@ -38,8 +39,6 @@ import { WorkspaceExperiments } from '../../../../experiments/workspace'
 import { EXPERIMENT_WORKSPACE_ID } from '../../../../cli/dvc/contract'
 import { copyOriginalColors } from '../../../../experiments/model/status/colors'
 import { Revision } from '../../../../plots/webview/contract'
-import { BaseWebview } from '../../../../webview'
-import { Experiment } from '../../../../experiments/webview/contract'
 
 suite('Experiments Tree Test Suite', () => {
   const disposable = getTimeSafeDisposer()
@@ -63,17 +62,15 @@ suite('Experiments Tree Test Suite', () => {
     it('should be able to toggle whether an experiment is shown in the plots webview with dvc.views.experiments.toggleStatus', async () => {
       const mockNow = getMockNow()
 
-      const { plots, messageSpy, plotsModel, experimentsModel } =
-        await buildPlots({ disposer: disposable })
-      messageSpy.restore()
-      const mockShow = stub(BaseWebview.prototype, 'show')
-
-      experimentsModel.setSelected([
-        { id: EXPERIMENT_WORKSPACE_ID },
-        { id: 'main' },
-        { id: 'test-branch' },
-        { id: 'exp-f13bca' }
-      ] as Experiment[])
+      const { messageSpy, plotsModel } = await buildPlotsWebview({
+        disposer: disposable,
+        selectedExperiments: [
+          EXPERIMENT_WORKSPACE_ID,
+          'main',
+          'test-branch',
+          'exp-f13bca'
+        ]
+      })
 
       const expectedRevisions: { displayColor: string; id: string }[] = []
 
@@ -86,13 +83,9 @@ suite('Experiments Tree Test Suite', () => {
         expectedRevisions.push({ displayColor, id })
       }
 
-      const webview = await plots.showWebview()
-
-      await webview.isReady()
-
       let updateCall = 1
       while (expectedRevisions.length > 0) {
-        const { selectedRevisions } = mockShow.lastCall.firstArg
+        const { selectedRevisions } = messageSpy.lastCall.firstArg
 
         expect(
           (selectedRevisions as Revision[]).map(({ displayColor, id }) => ({
@@ -101,17 +94,11 @@ suite('Experiments Tree Test Suite', () => {
           })),
           'a message is sent with colors for the currently selected experiments'
         ).to.deep.equal(expectedRevisions)
-        mockShow.resetHistory()
-        mockShow.resetBehavior()
+        messageSpy.resetHistory()
 
         const { id } = expectedRevisions.pop() as { id: string }
 
-        const messageSent = new Promise(resolve =>
-          mockShow.callsFake(() => {
-            resolve(undefined)
-            return Promise.resolve(true)
-          })
-        )
+        const messageSent = waitForSpyCall(messageSpy, messageSpy.callCount)
 
         bypassProcessManagerDebounce(mockNow, updateCall)
         const unSelected = await commands.executeCommand(
@@ -128,19 +115,14 @@ suite('Experiments Tree Test Suite', () => {
       }
 
       expect(
-        mockShow,
+        messageSpy,
         "when there are no experiments selected we don't send any template plots"
       ).to.be.calledWithMatch({
         template: null
       })
-      mockShow.resetHistory()
+      messageSpy.resetHistory()
 
-      const messageSent = new Promise(resolve =>
-        mockShow.callsFake(() => {
-          resolve(undefined)
-          return Promise.resolve(true)
-        })
-      )
+      const messageSent = waitForSpyCall(messageSpy, messageSpy.callCount)
 
       bypassProcessManagerDebounce(mockNow, updateCall)
       const selected = await commands.executeCommand(
@@ -191,7 +173,7 @@ suite('Experiments Tree Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to select / de-select experiments using dvc.views.experimentsTree.selectExperiments', async () => {
-      const { plots, plotsModel, messageSpy } = await buildPlots({
+      const { plotsModel, messageSpy } = await buildPlotsWebview({
         disposer: disposable
       })
 
@@ -203,8 +185,6 @@ suite('Experiments Tree Test Suite', () => {
         value: { id: label }
       }
 
-      await plots.showWebview()
-
       const inputAccepted = disposable.track(new EventEmitter<void>())
       const mockEvent = disposable.track(new EventEmitter()).event
 
@@ -226,14 +206,16 @@ suite('Experiments Tree Test Suite', () => {
         })
       )
 
+      const messageSent = waitForSpyCall(messageSpy, messageSpy.callCount)
+
       const selectExperiments = commands.executeCommand(
         RegisteredCommands.EXPERIMENT_SELECT
       )
 
       await quickPickCreated
       mockQuickPick.selectedItems = [selectedItem]
       inputAccepted.fire()
-      await selectExperiments
+      await Promise.all([selectExperiments, messageSent])
 
       expect(mockCreateQuickPick).to.be.calledOnce
 

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

@@ -19,8 +19,8 @@ import {
   bypassProcessManagerDebounce,
   closeAllEditors,
   getFirstArgOfLastCall,
-  getMessageReceivedEmitter,
-  getMockNow
+  getMockNow,
+  waitForSpyCall
 } from '../util'
 import { dvcDemoPath } from '../../util'
 import {
@@ -39,7 +39,6 @@ import { MessageFromWebviewType } from '../../../webview/contract'
 import { reorderObjectList, uniqueValues } from '../../../util/array'
 import * as Telemetry from '../../../telemetry'
 import { EventName } from '../../../telemetry/constants'
-import { SelectedExperimentWithColor } from '../../../experiments/model'
 import { ErrorItem } from '../../../path/selection/tree'
 import { isErrorItem } from '../../../tree'
 import { RegisteredCommands } from '../../../commands/external'
@@ -82,8 +81,11 @@ suite('Plots Test Suite', () => {
       )
       mockPlotsDiff.resetHistory()
 
+      const messageSent = waitForSpyCall(messageSpy, messageSpy.callCount)
+
       const webview = await plots.showWebview()
-      await webview.isReady()
+
+      await Promise.all([webview.isReady(), messageSent])
 
       expect(mockPlotsDiff).not.to.be.called
       expect(managedUpdateSpy).not.to.be.called
@@ -740,32 +742,24 @@ suite('Plots Test Suite', () => {
 
     it('should handle a message to manually refresh plot revisions from the webview', async () => {
       const mockNow = getMockNow()
-      const { data, plots, mockPlotsDiff, messageSpy } = await buildPlots({
-        disposer: disposable,
-        plotsDiff: plotsDiffFixture
-      })
-
-      messageSpy.restore()
-
-      const webview = await plots.showWebview()
-      mockPlotsDiff.resetHistory()
-      const instanceMessageSpy = spy(webview, 'show')
-      await webview.isReady()
+      const { messageSpy, mockMessageReceived, mockPlotsDiff } =
+        await buildPlotsWebview({
+          disposer: disposable,
+          plotsDiff: plotsDiffFixture
+        })
+      messageSpy.resetHistory()
 
       const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
-      const mockMessageReceived = getMessageReceivedEmitter(webview)
 
-      const dataUpdateEvent = new Promise(resolve =>
-        data.onDidUpdate(() => resolve(undefined))
-      )
+      const messageSent = waitForSpyCall(messageSpy, messageSpy.callCount)
 
       bypassProcessManagerDebounce(mockNow)
 
       mockMessageReceived.fire({
         type: MessageFromWebviewType.REFRESH_REVISIONS
       })
 
-      await dataUpdateEvent
+      await messageSent
 
       expect(mockSendTelemetryEvent).to.be.calledOnce
       expect(mockSendTelemetryEvent).to.be.calledWith(EventName.PLOTS_REFRESH)
@@ -779,7 +773,7 @@ suite('Plots Test Suite', () => {
         REVISIONS[1]
       )
       expect(
-        instanceMessageSpy,
+        messageSpy,
         'should call the plots webview with cached data before refreshing with data from the CLI'
       ).to.be.calledTwice
     }).timeout(WEBVIEW_TEST_TIMEOUT)
@@ -891,19 +885,11 @@ suite('Plots Test Suite', () => {
     })
 
     it('should send the correct data to the webview for flexible plots', async () => {
-      const { experiments, plots, messageSpy, mockPlotsDiff } =
-        await buildPlots({
-          disposer: disposable,
-          plotsDiff: multiSourcePlotsDiffFixture
-        })
-
-      stub(experiments, 'getSelectedRevisions').returns([
-        { id: REVISIONS[0] },
-        { id: REVISIONS[1] }
-      ] as SelectedExperimentWithColor[])
-
-      const webview = await plots.showWebview()
-      await webview.isReady()
+      const { messageSpy, mockPlotsDiff } = await buildPlotsWebview({
+        disposer: disposable,
+        plotsDiff: multiSourcePlotsDiffFixture,
+        selectedExperiments: [REVISIONS[0], REVISIONS[1]]
+      })
 
       expect(mockPlotsDiff).to.be.called
 

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

@@ -151,18 +151,21 @@ export const buildPlotsWebview = async ({
   expShow = expShowFixtureWithoutErrors,
   gitLog = gitLogFixture,
   plotsDiff = undefined,
+  selectedExperiments,
   rowOrder
 }: {
   availableNbCommits?: { [branch: string]: number }
   disposer: Disposer
   expShow?: ExpShowOutput
   gitLog?: string
   plotsDiff?: PlotsOutput | undefined
+  selectedExperiments?: string[]
   rowOrder?: { branch: string; sha: string }[]
 }) => {
   const {
     data,
     experiments,
+    experimentsModel,
     pathsModel,
     plots,
     plotsModel,
@@ -176,6 +179,11 @@ export const buildPlotsWebview = async ({
     plotsDiff,
     rowOrder
   })
+  if (selectedExperiments) {
+    experimentsModel.setSelected(
+      selectedExperiments.map(id => ({ id }) as Experiment)
+    )
+  }
 
   const webview = await plots.showWebview()
   await webview.isReady()
@@ -189,6 +197,7 @@ export const buildPlotsWebview = async ({
   return {
     data,
     experiments,
+    experimentsModel,
     messageSpy: instanceMessageSpy,
     mockMessageReceived,
     mockPlotsDiff,