From 43c4d121fb7015690ba7e3d1344675c7f71cbc85 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <sroy3@users.noreply.github.com>
Date: Sat, 26 Nov 2022 18:53:55 -0500
Subject: [PATCH] Fix the number of available revisions for multi view plots
 (#2836)

* Fix the aspect ratio of multi view plots

* Revert to old aspect ratio formula

* Fix the number of available revisions for multi view plots

* Revert "Fix the aspect ratio of multi view plots"

This reverts commit 71cd9617539eceaddc03f6c95e1cecaf2629c744.

* patch flexible plot test

Co-authored-by: Matt Seddon <mattseddon@hotmail.com>
---
 extension/src/plots/model/collect.ts               | 14 ++++++++------
 .../src/test/fixtures/plotsDiff/multiSource.ts     |  6 +++---
 extension/src/test/suite/plots/index.test.ts       | 13 +++++++++----
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/extension/src/plots/model/collect.ts b/extension/src/plots/model/collect.ts
index 53d2a69f50..15d71e912e 100644
--- a/extension/src/plots/model/collect.ts
+++ b/extension/src/plots/model/collect.ts
@@ -542,7 +542,9 @@ const transformRevisionData = (
 ): { revisions: string[]; datapoints: unknown[] } => {
   const field = multiSourceEncodingUpdate.strokeDash?.field
   const isMultiSource = !!field
-
+  const availableRevisions = selectedRevisions.filter(rev =>
+    Object.keys(revisionData).includes(rev)
+  )
   const transformNeeded =
     isMultiSource && (isMultiView || isConcatenatedField(field))
 
@@ -551,7 +553,7 @@ const transformRevisionData = (
       datapoints: selectedRevisions
         .flatMap(revision => revisionData?.[revision]?.[path])
         .filter(Boolean),
-      revisions: selectedRevisions
+      revisions: availableRevisions
     }
   }
 
@@ -562,12 +564,12 @@ const transformRevisionData = (
       datapoints: updateDatapoints(
         path,
         revisionData,
-        selectedRevisions,
+        availableRevisions,
         'rev',
         fields
       ),
       revisions: updateRevisions(
-        selectedRevisions,
+        availableRevisions,
         multiSourceEncodingUpdate.strokeDash.scale.domain
       )
     }
@@ -577,11 +579,11 @@ const transformRevisionData = (
     datapoints: updateDatapoints(
       path,
       revisionData,
-      selectedRevisions,
+      availableRevisions,
       field,
       fields
     ),
-    revisions: selectedRevisions
+    revisions: availableRevisions
   }
 }
 
diff --git a/extension/src/test/fixtures/plotsDiff/multiSource.ts b/extension/src/test/fixtures/plotsDiff/multiSource.ts
index 552e587c5e..d4081a0fa7 100644
--- a/extension/src/test/fixtures/plotsDiff/multiSource.ts
+++ b/extension/src/test/fixtures/plotsDiff/multiSource.ts
@@ -9660,7 +9660,7 @@ const data = {
             }
           }
         ],
-        HEAD: [
+        main: [
           {
             fpr: 0.0,
             tpr: 0.0,
@@ -38250,7 +38250,7 @@ const data = {
             }
           }
         ],
-        HEAD: [
+        main: [
           {
             precision: 0.36728395061728397,
             recall: 1.0,
@@ -197286,7 +197286,7 @@ const data = {
             }
           }
         ],
-        HEAD: [
+        main: [
           {
             actual: '0',
             predicted: '0',
diff --git a/extension/src/test/suite/plots/index.test.ts b/extension/src/test/suite/plots/index.test.ts
index 93f4b7eb0a..00f14ec6c9 100644
--- a/extension/src/test/suite/plots/index.test.ts
+++ b/extension/src/test/suite/plots/index.test.ts
@@ -39,6 +39,7 @@ import * as Telemetry from '../../../telemetry'
 import { EventName } from '../../../telemetry/constants'
 import { ExperimentStatus } from '../../../cli/dvc/contract'
 import { SelectedExperimentWithColor } from '../../../experiments/model'
+import { Experiments } from '../../../experiments'
 
 suite('Plots Test Suite', () => {
   const disposable = Disposable.fn()
@@ -763,14 +764,18 @@ suite('Plots Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should send the correct data to the webview for flexible plots', async () => {
-      const { plots, messageSpy, mockPlotsDiff, experiments } =
-        await buildPlots(disposable, multiSourcePlotsDiffFixture)
-
-      stub(experiments, 'getSelectedRevisions').returns([
+      stub(Experiments.prototype, 'getSelectedRevisions').returns([
         { label: 'workspace' },
         { label: 'main' }
       ] as SelectedExperimentWithColor[])
 
+      stub(Experiments.prototype, 'getBranchRevisions').returns([])
+
+      const { plots, messageSpy, mockPlotsDiff } = await buildPlots(
+        disposable,
+        multiSourcePlotsDiffFixture
+      )
+
       const webview = await plots.showWebview()
       await webview.isReady()