From 3d8404e921b03e08c0fa51ddeff088bbf3fe7057 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 27 Jul 2023 17:17:56 -0500
Subject: [PATCH 01/23] Add back changes in
 76ee2983d5276e18ab790b969fbccd2822447830

---
 .../comparisonTable/ComparisonTableRows.tsx   |  7 ++-
 .../comparisonTable/comparisonTableSlice.ts   | 23 +++++++-
 .../comparisonTable/styles.module.scss        | 54 ++++++++++++-------
 3 files changed, 61 insertions(+), 23 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
index aebc2bee21..9198e7a7d6 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
@@ -1,12 +1,13 @@
 import { ComparisonPlots } from 'dvc/src/plots/webview/contract'
 import React, { createRef, useEffect, useState } from 'react'
-import { useDispatch } from 'react-redux'
+import { useDispatch, useSelector } from 'react-redux'
 import { ComparisonTableColumn } from './ComparisonTableHead'
 import { ComparisonTableRow } from './ComparisonTableRow'
 import { changeRowHeight, DEFAULT_ROW_HEIGHT } from './comparisonTableSlice'
 import { RowDropTarget } from './RowDropTarget'
 import { DragDropContainer } from '../../../shared/components/dragDrop/DragDropContainer'
 import { reorderComparisonRows } from '../../util/messages'
+import { PlotsState } from '../../store'
 
 interface ComparisonTableRowsProps {
   plots: ComparisonPlots
@@ -22,6 +23,9 @@ export const ComparisionTableRows: React.FC<ComparisonTableRowsProps> = ({
   const [rowsOrder, setRowsOrder] = useState<string[]>([])
   const dispatch = useDispatch()
   const firstRowRef = createRef<HTMLTableSectionElement>()
+  const disabledDragPlotIds = useSelector(
+    (state: PlotsState) => state.comparison.disabledDragPlotIds
+  )
 
   useEffect(() => {
     setRowsOrder(plots.map(({ path }) => path))
@@ -67,6 +71,7 @@ export const ComparisionTableRows: React.FC<ComparisonTableRowsProps> = ({
       group="comparison-table"
       dropTarget={<RowDropTarget colSpan={columns.length} />}
       onLayoutChange={onLayoutChange}
+      disabledDropIds={disabledDragPlotIds}
       vertical
     />
   )
diff --git a/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts b/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts
index d9f9a3a1ed..930e85ba59 100644
--- a/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts
+++ b/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts
@@ -11,14 +11,18 @@ export interface ComparisonTableState extends PlotsComparisonData {
   isCollapsed: boolean
   hasData: boolean
   rowHeight: number
+  multiPlotValues: { [path: string]: number }
+  disabledDragPlotIds: string[]
 }
 
 export const DEFAULT_ROW_HEIGHT = 200
 
 export const comparisonTableInitialState: ComparisonTableState = {
+  disabledDragPlotIds: [],
   hasData: false,
   height: DEFAULT_HEIGHT[PlotsSection.COMPARISON_TABLE],
   isCollapsed: DEFAULT_SECTION_COLLAPSED[PlotsSection.COMPARISON_TABLE],
+  multiPlotValues: {},
   plots: [],
   revisions: [],
   rowHeight: DEFAULT_ROW_HEIGHT,
@@ -30,6 +34,9 @@ export const comparisonTableSlice = createSlice({
   initialState: comparisonTableInitialState,
   name: 'comparison',
   reducers: {
+    changeDisabledDragIds: (state, action: PayloadAction<string[]>) => {
+      state.disabledDragPlotIds = action.payload
+    },
     changeRowHeight: (state, action: PayloadAction<number>) => {
       state.rowHeight = action.payload
     },
@@ -42,6 +49,12 @@ export const comparisonTableSlice = createSlice({
     setCollapsed: (state, action: PayloadAction<boolean>) => {
       state.isCollapsed = action.payload
     },
+    setMultiPlotValue: (
+      state,
+      action: PayloadAction<{ path: string; value: number }>
+    ) => {
+      state.multiPlotValues[action.payload.path] = action.payload.value
+    },
     update: (state, action: PayloadAction<PlotsComparisonData>) => {
       if (!action.payload) {
         return comparisonTableInitialState
@@ -55,7 +68,13 @@ export const comparisonTableSlice = createSlice({
   }
 })
 
-export const { update, setCollapsed, changeSize, changeRowHeight } =
-  comparisonTableSlice.actions
+export const {
+  update,
+  setCollapsed,
+  setMultiPlotValue,
+  changeSize,
+  changeDisabledDragIds,
+  changeRowHeight
+} = comparisonTableSlice.actions
 
 export default comparisonTableSlice.reducer
diff --git a/webview/src/plots/components/comparisonTable/styles.module.scss b/webview/src/plots/components/comparisonTable/styles.module.scss
index 9784d666aa..2139e59e63 100644
--- a/webview/src/plots/components/comparisonTable/styles.module.scss
+++ b/webview/src/plots/components/comparisonTable/styles.module.scss
@@ -88,25 +88,6 @@ $gap: 4px;
   transform: rotate(0deg);
 }
 
-.imageWrapper {
-  width: 100%;
-  display: block;
-  padding: 0;
-  border: 0;
-}
-
-.noImage {
-  background-color: $bg-color;
-  border-style: solid;
-  border-width: thin;
-  border-color: $bg-color;
-}
-
-.noImageContent {
-  padding-top: 25%;
-  padding-bottom: 25%;
-}
-
 .rowToggler {
   border: none;
   background: none;
@@ -170,12 +151,45 @@ $gap: 4px;
   max-height: 0;
 }
 
-.cell img {
+.image {
   width: 100%;
   height: auto;
   vertical-align: middle;
 }
 
+.imageWrapper {
+  width: 100%;
+  display: block;
+  padding: 0;
+  border: 0;
+}
+
+.noImageContent {
+  padding-top: 25%;
+  padding-bottom: 25%;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+}
+
+.multiImageWrapper {
+  .image,
+  .noImageContent {
+    height: 380px;
+    object-fit: contain;
+  }
+}
+
+.multiImageSlider {
+  display: flex;
+  height: 40px;
+  align-items: center;
+  justify-content: center;
+  color: $fg-color;
+  background-color: $bg-color;
+  box-shadow: inset 40px 40px $fg-transparency-1;
+}
+
 .experimentName {
   color: $meta-cell-color;
 }

From ffd42f816807bd8ce01fe6f41b04a20ced19551d Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 27 Jul 2023 17:18:53 -0500
Subject: [PATCH 02/23] Revert "Simplify first pr changes"

This reverts commit 739d752f298dac8bfa54195c07f04da79512fbf7.
---
 extension/src/plots/model/collect.test.ts     |   5 +-
 extension/src/plots/model/collect.ts          |   5 +-
 extension/src/plots/paths/model.test.ts       |   8 +-
 extension/src/plots/paths/model.ts            |  39 ++++++-
 .../fixtures/plotsDiff/comparison/multi.ts    |   5 -
 .../src/test/fixtures/plotsDiff/index.ts      |  41 ++++---
 mockDvc/.dvc/.gitignore                       |   3 +
 .../comparisonTable/ComparisonTable.test.tsx  |   4 +
 .../comparisonTable/ComparisonTableRow.tsx    |   7 +-
 .../cell/ComparisonTableCell.tsx              |   4 +-
 .../cell/ComparisonTableMultiCell.tsx         | 100 ++++++++++++++++++
 .../src/stories/ComparisonTable.stories.tsx   |   6 --
 12 files changed, 183 insertions(+), 44 deletions(-)
 delete mode 100644 extension/src/test/fixtures/plotsDiff/comparison/multi.ts
 create mode 100644 mockDvc/.dvc/.gitignore
 create mode 100644 webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx

diff --git a/extension/src/plots/model/collect.test.ts b/extension/src/plots/model/collect.test.ts
index 566e3657d8..ff91edec96 100644
--- a/extension/src/plots/model/collect.test.ts
+++ b/extension/src/plots/model/collect.test.ts
@@ -22,7 +22,7 @@ import {
   TemplatePlot
 } from '../webview/contract'
 import { exists } from '../../fileSystem'
-import { REVISIONS } from '../../test/fixtures/plotsDiff'
+import { REVISIONS, getMultiImagePaths } from '../../test/fixtures/plotsDiff'
 
 const mockedExists = jest.mocked(exists)
 
@@ -106,7 +106,8 @@ describe('collectTemplates', () => {
     expect(Object.keys(templates)).toStrictEqual([
       logsLossPath,
       join('logs', 'acc.tsv'),
-      'predictions.json'
+      'predictions.json',
+      ...getMultiImagePaths()
     ])
 
     expect(JSON.parse(templates[logsLossPath])).toStrictEqual(content)
diff --git a/extension/src/plots/model/collect.ts b/extension/src/plots/model/collect.ts
index c728edc2cf..98da7b27a3 100644
--- a/extension/src/plots/model/collect.ts
+++ b/extension/src/plots/model/collect.ts
@@ -36,6 +36,7 @@ import {
 import { StrokeDashEncoding } from '../multiSource/constants'
 import { exists } from '../../fileSystem'
 import { hasKey } from '../../util/object'
+import { getParent, getPathArray } from '../../fileSystem/util'
 import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
 
 export const getCustomPlotId = (metric: string, param: string) =>
@@ -262,7 +263,9 @@ const collectSelectedPathComparisonPlots = ({
   getComparisonPlotImg: (id: string, path: string) => ComparisonPlotImg
 }) => {
   const isMulti = MULTI_IMAGE_PATH_REG.test(path)
-  const pathLabel = path
+  const pathLabel = isMulti
+    ? (getParent(getPathArray(path), 0) as string)
+    : path
   const doesPathExist = acc[pathLabel]
 
   if (!doesPathExist) {
diff --git a/extension/src/plots/paths/model.test.ts b/extension/src/plots/paths/model.test.ts
index fb83d660a7..1c3ec47250 100644
--- a/extension/src/plots/paths/model.test.ts
+++ b/extension/src/plots/paths/model.test.ts
@@ -6,7 +6,7 @@ import { buildMockMemento } from '../../test/util'
 import { PlotsType, TemplatePlotGroup } from '../webview/contract'
 import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 import { ErrorsModel } from '../errors/model'
-import { REVISIONS } from '../../test/fixtures/plotsDiff'
+import { REVISIONS, getMultiImagePaths } from '../../test/fixtures/plotsDiff'
 
 describe('PathsModel', () => {
   const mockDvcRoot = 'test'
@@ -340,13 +340,15 @@ describe('PathsModel', () => {
     expect(model.getComparisonPaths()).toStrictEqual([
       join('plots', 'acc.png'),
       join('plots', 'heatmap.png'),
-      join('plots', 'loss.png')
+      join('plots', 'loss.png'),
+      ...getMultiImagePaths()
     ])
 
     const newOrder = [
       join('plots', 'heatmap.png'),
       join('plots', 'acc.png'),
-      join('plots', 'loss.png')
+      join('plots', 'loss.png'),
+      ...getMultiImagePaths()
     ]
 
     model.setComparisonPathsOrder(newOrder)
diff --git a/extension/src/plots/paths/model.ts b/extension/src/plots/paths/model.ts
index dd8aa3ea7a..f682a0c8f7 100644
--- a/extension/src/plots/paths/model.ts
+++ b/extension/src/plots/paths/model.ts
@@ -18,6 +18,7 @@ import { isDvcError } from '../../cli/dvc/reader'
 import { PlotsOutputOrError } from '../../cli/dvc/contract'
 import { getErrorTooltip } from '../../tree'
 import { ErrorsModel } from '../errors/model'
+import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
 
 export class PathsModel extends PathSelectionModel<PlotPath> {
   private readonly errors: ErrorsModel
@@ -126,10 +127,18 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
   }
 
   public getComparisonPaths() {
-    return performSimpleOrderedUpdate(
-      this.comparisonPathsOrder,
+    // TBD do we need to do sorting everytime we need plots
+    // should we save sorted/grouped data into its own variables to use?
+    const { multiImagePaths, paths } = this.sortComparisonPaths(
       this.getPathsByType(PathType.COMPARISON)
     )
+
+    return performSimpleOrderedUpdate(
+      this.comparisonPathsOrder.flatMap(path =>
+        path.endsWith('image') ? multiImagePaths : path
+      ),
+      paths
+    )
   }
 
   public setComparisonPathsOrder(order: string[]) {
@@ -162,6 +171,32 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
       .map(({ path }) => path)
   }
 
+  private sortComparisonPaths(originalPaths: string[]) {
+    const multiImagePaths: string[] = []
+    const paths: string[] = []
+    let multiImageInd
+
+    for (const [ind, path] of originalPaths.entries()) {
+      const isMulti = MULTI_IMAGE_PATH_REG.test(path)
+      if (!isMulti) {
+        paths.push(path)
+        continue
+      }
+      if (!multiImageInd) {
+        multiImageInd = ind
+      }
+      multiImagePaths.push(path)
+    }
+
+    multiImagePaths.sort((path1, path2) =>
+      path1.localeCompare(path2, undefined, { numeric: true })
+    )
+
+    paths.splice(multiImageInd || 0, 0, ...multiImagePaths)
+
+    return { multiImagePaths, paths }
+  }
+
   private filterChildren(path: string | undefined): PlotPath[] {
     return this.data.filter(element => {
       if (!this.hasRevisions(element)) {
diff --git a/extension/src/test/fixtures/plotsDiff/comparison/multi.ts b/extension/src/test/fixtures/plotsDiff/comparison/multi.ts
deleted file mode 100644
index 8c8a2bd1ab..0000000000
--- a/extension/src/test/fixtures/plotsDiff/comparison/multi.ts
+++ /dev/null
@@ -1,5 +0,0 @@
-import { getComparisonWebviewMessage } from '..'
-
-const data = getComparisonWebviewMessage('.', undefined, true)
-
-export default data
diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index be557941d8..633a67df0b 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -363,7 +363,8 @@ const multipleVega = (length: number) => {
   }
   return plots
 }
-
+// TBD I don't think we need to add muliImageData to the main fixture used in tests across the project
+// we should put multi image data into its own fixture and use in desired frontend/backend tests
 const getMultiImageData = (
   baseUrl: string,
   joinFunc: (...segments: string[]) => string,
@@ -472,11 +473,7 @@ const getImageData = (baseUrl: string, joinFunc = join) => ({
       revisions: ['exp-83425'],
       url: joinFunc(baseUrl, '1ba7bcd_plots_loss.png')
     }
-  ]
-})
-
-const getImageDataWithMultiImgs = (baseUrl: string, joinFunc = join) => ({
-  ...getImageData(baseUrl, joinFunc),
+  ],
   ...getMultiImageData(baseUrl, joinFunc, [
     EXPERIMENT_WORKSPACE_ID,
     'main',
@@ -486,6 +483,14 @@ const getImageDataWithMultiImgs = (baseUrl: string, joinFunc = join) => ({
   ])
 })
 
+export const getMultiImagePaths = () => {
+  const paths = []
+  for (let i = 0; i < 15; i++) {
+    paths.push(join('plots', 'image', `${i}.jpg`))
+  }
+  return paths
+}
+
 export const getOutput = (baseUrl: string): PlotsOutput => ({
   data: {
     ...getImageData(baseUrl),
@@ -797,25 +802,20 @@ export const MOCK_IMAGE_MTIME = 946684800000
 
 export const getComparisonWebviewMessage = (
   baseUrl: string,
-  joinFunc: (...args: string[]) => string = join,
-  addMulti?: boolean
+  joinFunc: (...args: string[]) => string = join
 ): PlotsComparisonData => {
   const plotAcc: {
     [path: string]: { path: string; revisions: ComparisonRevisionData }
   } = {}
 
-  for (const [path, plots] of Object.entries(
-    addMulti
-      ? getImageDataWithMultiImgs(baseUrl, joinFunc)
-      : getImageData(baseUrl, joinFunc)
-  )) {
+  for (const [path, plots] of Object.entries(getImageData(baseUrl, joinFunc))) {
     const multiImagePath = joinFunc('plots', 'image')
     const isMulti = path.includes(multiImagePath)
-    const pathLabel = path
+    const pathName = isMulti ? multiImagePath : path
 
-    if (!plotAcc[pathLabel]) {
-      plotAcc[pathLabel] = {
-        path: pathLabel,
+    if (!plotAcc[pathName]) {
+      plotAcc[pathName] = {
+        path: pathName,
         revisions: {}
       }
     }
@@ -826,15 +826,14 @@ export const getComparisonWebviewMessage = (
         continue
       }
 
-      if (!plotAcc[pathLabel].revisions[id]) {
-        plotAcc[pathLabel].revisions[id] = {
+      if (!plotAcc[pathName].revisions[id]) {
+        plotAcc[pathName].revisions[id] = {
           id,
           imgOrImgs: [],
           isMulti
         }
       }
-
-      plotAcc[pathLabel].revisions[id].imgOrImgs.push({
+      plotAcc[pathName].revisions[id].imgOrImgs.push({
         url: `${url}?${MOCK_IMAGE_MTIME}`,
         errors: undefined,
         loading: false
diff --git a/mockDvc/.dvc/.gitignore b/mockDvc/.dvc/.gitignore
new file mode 100644
index 0000000000..528f30c71c
--- /dev/null
+++ b/mockDvc/.dvc/.gitignore
@@ -0,0 +1,3 @@
+/config.local
+/tmp
+/cache
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
index dc4305ee84..f890cdf77f 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
@@ -280,7 +280,9 @@ describe('ComparisonTable', () => {
             imgOrImgs: [
               {
                 errors: undefined,
+                id: revisionWithNoData,
                 loading: false,
+                path: '',
                 url: undefined
               }
             ],
@@ -318,7 +320,9 @@ describe('ComparisonTable', () => {
             imgOrImgs: [
               {
                 errors: ['this is an error'],
+                id: revisionWithNoData,
                 loading: false,
+                path: '',
                 url: undefined
               }
             ],
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx
index f92be3e0b2..c6259b596f 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx
@@ -4,6 +4,7 @@ import cx from 'classnames'
 import { useSelector } from 'react-redux'
 import styles from './styles.module.scss'
 import { ComparisonTableCell } from './cell/ComparisonTableCell'
+import { ComparisonTableMultiCell } from './cell/ComparisonTableMultiCell'
 import { Icon } from '../../../shared/components/Icon'
 import { ChevronDown, ChevronRight } from '../../../shared/components/icons'
 import { PlotsState } from '../../store'
@@ -76,7 +77,11 @@ export const ComparisonTableRow: React.FC<ComparisonTableRowProps> = ({
                 data-testid="row-images"
                 className={cx(styles.cell, { [styles.cellHidden]: !isShown })}
               >
-                <ComparisonTableCell plot={plot} path={path} />
+                {plot.isMulti ? (
+                  <ComparisonTableMultiCell plot={plot} path={path} />
+                ) : (
+                  <ComparisonTableCell plot={plot} path={path} />
+                )}
               </div>
             </td>
           )
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
index d2075b3aea..bca1ffad4f 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
@@ -9,9 +9,7 @@ export const ComparisonTableCell: React.FC<{
   path: string
   plot: ComparisonPlot
 }> = ({ path, plot }) => {
-  const plotImg = plot.imgOrImgs
-    ? plot.imgOrImgs[0]
-    : { errors: undefined, loading: false, url: undefined }
+  const plotImg = plot.imgOrImgs[0]
   const loading = plotImg.loading
   const missing = !loading && !plotImg.url
 
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
new file mode 100644
index 0000000000..4a0e69e6b4
--- /dev/null
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
@@ -0,0 +1,100 @@
+import React, { useCallback, MouseEvent, KeyboardEvent } from 'react'
+import cx from 'classnames'
+import { useDispatch, useSelector } from 'react-redux'
+import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
+import {
+  changeDisabledDragIds,
+  setMultiPlotValue
+} from '../comparisonTableSlice'
+import { PlotsState } from '../../../store'
+import styles from '../styles.module.scss'
+import { ComparisonTableLoadingCell } from './ComparisonTableLoadingCell'
+import { ComparisonTableMissingCell } from './ComparisonTableMissingCell'
+import { zoomPlot } from '../../../util/messages'
+
+export const ComparisonTableMultiCell: React.FC<{
+  path: string
+  plot: ComparisonPlot
+}> = ({ path, plot }) => {
+  const multiValues = useSelector(
+    (state: PlotsState) => state.comparison.multiPlotValues
+  )
+  const dispatch = useDispatch()
+  const currentStep = multiValues[path] || 0
+
+  const { loading, url } = plot.imgOrImgs[currentStep]
+  const missing = !loading && !url
+
+  const addDisabled = useCallback(() => {
+    dispatch(changeDisabledDragIds([path]))
+  }, [dispatch, path])
+
+  const removeDisabled = useCallback(() => {
+    dispatch(changeDisabledDragIds([]))
+  }, [dispatch])
+
+  const disableClick = useCallback((e: MouseEvent | KeyboardEvent) => {
+    e.stopPropagation()
+  }, [])
+
+  const slider = (
+    // The div element has children that allow keyboard interaction
+    // eslint-disable-next-line jsx-a11y/no-static-element-interactions
+    <div
+      className={styles.multiImageSlider}
+      onMouseEnter={addDisabled}
+      onMouseLeave={removeDisabled}
+      onClick={disableClick}
+      onKeyDown={disableClick}
+    >
+      <label htmlFor={`${plot.id}-step`}>Step</label>
+      <input
+        name={`${plot.id}-step`}
+        min="0"
+        max={plot.imgOrImgs.length - 1}
+        value={currentStep}
+        type="range"
+        onChange={event => {
+          dispatch(
+            setMultiPlotValue({ path, value: Number(event.target.value) })
+          )
+        }}
+      />
+      <p>{currentStep}</p>
+    </div>
+  )
+
+  if (loading) {
+    return (
+      <div className={styles.multiImageWrapper}>
+        <ComparisonTableLoadingCell />
+        {slider}
+      </div>
+    )
+  }
+
+  if (missing) {
+    return (
+      <div className={styles.multiImageWrapper}>
+        <ComparisonTableMissingCell plot={plot.imgOrImgs[currentStep]} />
+        {slider}
+      </div>
+    )
+  }
+
+  return (
+    <button
+      className={cx(styles.imageWrapper, styles.multiImageWrapper)}
+      onClick={() => url && zoomPlot(url)}
+      data-testid="image-plot-button"
+    >
+      <img
+        className={styles.image}
+        draggable={false}
+        src={url}
+        alt={`${currentStep} of ${path} (${plot.id})`}
+      />
+      {slider}
+    </button>
+  )
+}
diff --git a/webview/src/stories/ComparisonTable.stories.tsx b/webview/src/stories/ComparisonTable.stories.tsx
index 81808f765e..c3b5ab3e7f 100644
--- a/webview/src/stories/ComparisonTable.stories.tsx
+++ b/webview/src/stories/ComparisonTable.stories.tsx
@@ -11,7 +11,6 @@ import {
   PlotsComparisonData
 } from 'dvc/src/plots/webview/contract'
 import comparisonTableFixture from 'dvc/src/test/fixtures/plotsDiff/comparison'
-import comparisonTableMultiFixture from 'dvc/src/test/fixtures/plotsDiff/comparison/multi'
 import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
 import { DISABLE_CHROMATIC_SNAPSHOTS } from './util'
 import { ComparisonTable } from '../plots/components/comparisonTable/ComparisonTable'
@@ -104,11 +103,6 @@ const removeImages = (
   return filteredRevisionData
 }
 
-export const WithMultiImages = Template.bind({})
-WithMultiImages.args = {
-  plots: comparisonTableMultiFixture.plots
-}
-
 export const WithMissingData = Template.bind({})
 WithMissingData.args = {
   plots: comparisonTableFixture.plots.map(({ path, revisions }) => ({

From 4af98dd18f0e20f72b4473b0222ab6da3efefb60 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 27 Jul 2023 17:20:31 -0500
Subject: [PATCH 03/23] Delete .gitignore

---
 mockDvc/.dvc/.gitignore | 3 ---
 1 file changed, 3 deletions(-)
 delete mode 100644 mockDvc/.dvc/.gitignore

diff --git a/mockDvc/.dvc/.gitignore b/mockDvc/.dvc/.gitignore
deleted file mode 100644
index 528f30c71c..0000000000
--- a/mockDvc/.dvc/.gitignore
+++ /dev/null
@@ -1,3 +0,0 @@
-/config.local
-/tmp
-/cache

From 119b63a6259fc40a783de4113c3a6f87ddde7690 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 27 Jul 2023 17:29:14 -0500
Subject: [PATCH 04/23] fix revert issues

---
 extension/src/plots/model/collect.test.ts     |  5 +--
 extension/src/plots/paths/model.test.ts       |  8 ++--
 .../fixtures/plotsDiff/comparison/multi.ts    |  5 +++
 .../src/test/fixtures/plotsDiff/index.ts      | 41 ++++++++++---------
 .../cell/ComparisonTableMultiCell.tsx         |  6 +--
 5 files changed, 34 insertions(+), 31 deletions(-)
 create mode 100644 extension/src/test/fixtures/plotsDiff/comparison/multi.ts

diff --git a/extension/src/plots/model/collect.test.ts b/extension/src/plots/model/collect.test.ts
index ff91edec96..566e3657d8 100644
--- a/extension/src/plots/model/collect.test.ts
+++ b/extension/src/plots/model/collect.test.ts
@@ -22,7 +22,7 @@ import {
   TemplatePlot
 } from '../webview/contract'
 import { exists } from '../../fileSystem'
-import { REVISIONS, getMultiImagePaths } from '../../test/fixtures/plotsDiff'
+import { REVISIONS } from '../../test/fixtures/plotsDiff'
 
 const mockedExists = jest.mocked(exists)
 
@@ -106,8 +106,7 @@ describe('collectTemplates', () => {
     expect(Object.keys(templates)).toStrictEqual([
       logsLossPath,
       join('logs', 'acc.tsv'),
-      'predictions.json',
-      ...getMultiImagePaths()
+      'predictions.json'
     ])
 
     expect(JSON.parse(templates[logsLossPath])).toStrictEqual(content)
diff --git a/extension/src/plots/paths/model.test.ts b/extension/src/plots/paths/model.test.ts
index 1c3ec47250..fb83d660a7 100644
--- a/extension/src/plots/paths/model.test.ts
+++ b/extension/src/plots/paths/model.test.ts
@@ -6,7 +6,7 @@ import { buildMockMemento } from '../../test/util'
 import { PlotsType, TemplatePlotGroup } from '../webview/contract'
 import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 import { ErrorsModel } from '../errors/model'
-import { REVISIONS, getMultiImagePaths } from '../../test/fixtures/plotsDiff'
+import { REVISIONS } from '../../test/fixtures/plotsDiff'
 
 describe('PathsModel', () => {
   const mockDvcRoot = 'test'
@@ -340,15 +340,13 @@ describe('PathsModel', () => {
     expect(model.getComparisonPaths()).toStrictEqual([
       join('plots', 'acc.png'),
       join('plots', 'heatmap.png'),
-      join('plots', 'loss.png'),
-      ...getMultiImagePaths()
+      join('plots', 'loss.png')
     ])
 
     const newOrder = [
       join('plots', 'heatmap.png'),
       join('plots', 'acc.png'),
-      join('plots', 'loss.png'),
-      ...getMultiImagePaths()
+      join('plots', 'loss.png')
     ]
 
     model.setComparisonPathsOrder(newOrder)
diff --git a/extension/src/test/fixtures/plotsDiff/comparison/multi.ts b/extension/src/test/fixtures/plotsDiff/comparison/multi.ts
new file mode 100644
index 0000000000..8c8a2bd1ab
--- /dev/null
+++ b/extension/src/test/fixtures/plotsDiff/comparison/multi.ts
@@ -0,0 +1,5 @@
+import { getComparisonWebviewMessage } from '..'
+
+const data = getComparisonWebviewMessage('.', undefined, true)
+
+export default data
diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index 633a67df0b..be557941d8 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -363,8 +363,7 @@ const multipleVega = (length: number) => {
   }
   return plots
 }
-// TBD I don't think we need to add muliImageData to the main fixture used in tests across the project
-// we should put multi image data into its own fixture and use in desired frontend/backend tests
+
 const getMultiImageData = (
   baseUrl: string,
   joinFunc: (...segments: string[]) => string,
@@ -473,7 +472,11 @@ const getImageData = (baseUrl: string, joinFunc = join) => ({
       revisions: ['exp-83425'],
       url: joinFunc(baseUrl, '1ba7bcd_plots_loss.png')
     }
-  ],
+  ]
+})
+
+const getImageDataWithMultiImgs = (baseUrl: string, joinFunc = join) => ({
+  ...getImageData(baseUrl, joinFunc),
   ...getMultiImageData(baseUrl, joinFunc, [
     EXPERIMENT_WORKSPACE_ID,
     'main',
@@ -483,14 +486,6 @@ const getImageData = (baseUrl: string, joinFunc = join) => ({
   ])
 })
 
-export const getMultiImagePaths = () => {
-  const paths = []
-  for (let i = 0; i < 15; i++) {
-    paths.push(join('plots', 'image', `${i}.jpg`))
-  }
-  return paths
-}
-
 export const getOutput = (baseUrl: string): PlotsOutput => ({
   data: {
     ...getImageData(baseUrl),
@@ -802,20 +797,25 @@ export const MOCK_IMAGE_MTIME = 946684800000
 
 export const getComparisonWebviewMessage = (
   baseUrl: string,
-  joinFunc: (...args: string[]) => string = join
+  joinFunc: (...args: string[]) => string = join,
+  addMulti?: boolean
 ): PlotsComparisonData => {
   const plotAcc: {
     [path: string]: { path: string; revisions: ComparisonRevisionData }
   } = {}
 
-  for (const [path, plots] of Object.entries(getImageData(baseUrl, joinFunc))) {
+  for (const [path, plots] of Object.entries(
+    addMulti
+      ? getImageDataWithMultiImgs(baseUrl, joinFunc)
+      : getImageData(baseUrl, joinFunc)
+  )) {
     const multiImagePath = joinFunc('plots', 'image')
     const isMulti = path.includes(multiImagePath)
-    const pathName = isMulti ? multiImagePath : path
+    const pathLabel = path
 
-    if (!plotAcc[pathName]) {
-      plotAcc[pathName] = {
-        path: pathName,
+    if (!plotAcc[pathLabel]) {
+      plotAcc[pathLabel] = {
+        path: pathLabel,
         revisions: {}
       }
     }
@@ -826,14 +826,15 @@ export const getComparisonWebviewMessage = (
         continue
       }
 
-      if (!plotAcc[pathName].revisions[id]) {
-        plotAcc[pathName].revisions[id] = {
+      if (!plotAcc[pathLabel].revisions[id]) {
+        plotAcc[pathLabel].revisions[id] = {
           id,
           imgOrImgs: [],
           isMulti
         }
       }
-      plotAcc[pathName].revisions[id].imgOrImgs.push({
+
+      plotAcc[pathLabel].revisions[id].imgOrImgs.push({
         url: `${url}?${MOCK_IMAGE_MTIME}`,
         errors: undefined,
         loading: false
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
index 4a0e69e6b4..2e3801bfbd 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
@@ -2,14 +2,14 @@ import React, { useCallback, MouseEvent, KeyboardEvent } from 'react'
 import cx from 'classnames'
 import { useDispatch, useSelector } from 'react-redux'
 import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
+import { ComparisonTableLoadingCell } from './ComparisonTableLoadingCell'
+import { ComparisonTableMissingCell } from './ComparisonTableMissingCell'
+import styles from '../styles.module.scss'
 import {
   changeDisabledDragIds,
   setMultiPlotValue
 } from '../comparisonTableSlice'
 import { PlotsState } from '../../../store'
-import styles from '../styles.module.scss'
-import { ComparisonTableLoadingCell } from './ComparisonTableLoadingCell'
-import { ComparisonTableMissingCell } from './ComparisonTableMissingCell'
 import { zoomPlot } from '../../../util/messages'
 
 export const ComparisonTableMultiCell: React.FC<{

From b6114e8ff239d9edcbb2f57949e150117f1c75db Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 28 Jul 2023 11:42:00 -0500
Subject: [PATCH 05/23] Clean up paths model

---
 extension/src/plots/paths/model.ts | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/extension/src/plots/paths/model.ts b/extension/src/plots/paths/model.ts
index f682a0c8f7..d5bcb97601 100644
--- a/extension/src/plots/paths/model.ts
+++ b/extension/src/plots/paths/model.ts
@@ -127,22 +127,21 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
   }
 
   public getComparisonPaths() {
-    // TBD do we need to do sorting everytime we need plots
-    // should we save sorted/grouped data into its own variables to use?
-    const { multiImagePaths, paths } = this.sortComparisonPaths(
+    const { paths } = this.sortComparisonPaths(
       this.getPathsByType(PathType.COMPARISON)
     )
 
-    return performSimpleOrderedUpdate(
-      this.comparisonPathsOrder.flatMap(path =>
-        path.endsWith('image') ? multiImagePaths : path
-      ),
-      paths
-    )
+    return performSimpleOrderedUpdate(this.comparisonPathsOrder, paths)
   }
 
   public setComparisonPathsOrder(order: string[]) {
-    this.comparisonPathsOrder = order
+    const { multiImagePaths } = this.sortComparisonPaths(
+      this.getPathsByType(PathType.COMPARISON)
+    )
+
+    this.comparisonPathsOrder = order.flatMap(path =>
+      path.endsWith('image') ? multiImagePaths : path
+    )
     this.persist(
       PersistenceKey.PLOT_COMPARISON_PATHS_ORDER,
       this.comparisonPathsOrder

From 00d905b774f86bc7bc7e45d41a9a0e5b90c557ac Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 28 Jul 2023 12:04:35 -0500
Subject: [PATCH 06/23] Fix slider breaking for missing images

---
 .../cell/ComparisonTableMultiCell.tsx         | 86 +++++++++----------
 1 file changed, 39 insertions(+), 47 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
index 955488f6cd..0065907b8f 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
@@ -37,52 +37,7 @@ export const ComparisonTableMultiCell: React.FC<{
     e.stopPropagation()
   }, [])
 
-  const slider = (
-    // The div element has children that allow keyboard interaction
-    // eslint-disable-next-line jsx-a11y/no-static-element-interactions
-    <div
-      className={styles.multiImageSlider}
-      onMouseEnter={addDisabled}
-      onMouseLeave={removeDisabled}
-      onClick={disableClick}
-      onKeyDown={disableClick}
-    >
-      <label htmlFor={`${plot.id}-step`}>Step</label>
-      <input
-        name={`${plot.id}-step`}
-        min="0"
-        max={plot.imgs.length - 1}
-        value={currentStep}
-        type="range"
-        onChange={event => {
-          dispatch(
-            setMultiPlotValue({ path, value: Number(event.target.value) })
-          )
-        }}
-      />
-      <p>{currentStep}</p>
-    </div>
-  )
-
-  if (loading) {
-    return (
-      <div className={styles.multiImageWrapper}>
-        <ComparisonTableLoadingCell />
-        {slider}
-      </div>
-    )
-  }
-
-  if (missing) {
-    return (
-      <div className={styles.multiImageWrapper}>
-        <ComparisonTableMissingCell plot={plot.imgs[currentStep]} />
-        {slider}
-      </div>
-    )
-  }
-
-  return (
+  let imageContent = (
     <button
       className={cx(styles.imageWrapper, styles.multiImageWrapper)}
       onClick={() => url && zoomPlot(url)}
@@ -94,7 +49,44 @@ export const ComparisonTableMultiCell: React.FC<{
         src={url}
         alt={`${currentStep} of ${path} (${plot.id})`}
       />
-      {slider}
     </button>
   )
+
+  if (loading) {
+    imageContent = <ComparisonTableLoadingCell />
+  }
+
+  if (missing) {
+    imageContent = <ComparisonTableMissingCell plot={plot.imgs[currentStep]} />
+  }
+
+  return (
+    <div className={styles.multiImageWrapper}>
+      {imageContent}
+      {/* The div element has children that allow keyboard interaction */}
+      {/* eslint-disable-next-line jsx-a11y/no-static-element-interactions */}
+      <div
+        className={styles.multiImageSlider}
+        onMouseEnter={addDisabled}
+        onMouseLeave={removeDisabled}
+        onClick={disableClick}
+        onKeyDown={disableClick}
+      >
+        <label htmlFor={`${plot.id}-step`}>Step</label>
+        <input
+          name={`${plot.id}-step`}
+          min="0"
+          max={plot.imgs.length - 1}
+          value={currentStep}
+          type="range"
+          onChange={event => {
+            dispatch(
+              setMultiPlotValue({ path, value: Number(event.target.value) })
+            )
+          }}
+        />
+        <p>{currentStep}</p>
+      </div>
+    </div>
+  )
 }

From 8a60a60854ddd2cd728bea3b5fbce384a00395cb Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 28 Jul 2023 12:05:29 -0500
Subject: [PATCH 07/23] Fix jest tests

---
 .../components/comparisonTable/ComparisonTableRows.tsx      | 5 ++++-
 .../components/comparisonTable/cell/ComparisonTableCell.tsx | 6 ++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
index 9198e7a7d6..14a3826e9a 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
@@ -49,7 +49,10 @@ export const ComparisionTableRows: React.FC<ComparisonTableRowsProps> = ({
           path={path}
           plots={columns.map(column => ({
             ...revs[column.id],
-            id: column.id
+            id: column.id,
+            imgs: revs[column.id]?.imgs || [
+              { errors: undefined, loading: false, url: undefined }
+            ]
           }))}
           nbColumns={columns.length}
           pinnedColumn={pinnedColumn}
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
index 3ca44602ea..63158f62fc 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
@@ -9,10 +9,8 @@ export const ComparisonTableCell: React.FC<{
   path: string
   plot: ComparisonPlot
 }> = ({ path, plot }) => {
-  const plotImg = plot.imgs
-    ? plot.imgs[0]
-    : // TBD according to testing, imgs could be undefined. Need to review how.
-      { errors: undefined, loading: false, url: undefined }
+  const plotImg = plot.imgs[0]
+
   const loading = plotImg.loading
   const missing = !loading && !plotImg.url
 

From 6ef31f03b56018126ef04fe9626a16264f05dea7 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 28 Jul 2023 19:55:40 -0500
Subject: [PATCH 08/23] Add test for frontend

---
 .../src/test/fixtures/plotsDiff/index.ts      |   3 +-
 webview/src/plots/components/App.test.tsx     | 138 ++++++++++++++++++
 .../comparisonTable/ComparisonTableRows.tsx   |   7 +-
 .../cell/ComparisonTableMultiCell.tsx         |  14 +-
 4 files changed, 148 insertions(+), 14 deletions(-)

diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index 476cac6d85..1e363babf9 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -810,8 +810,7 @@ export const getComparisonWebviewMessage = (
       : getImageData(baseUrl, joinFunc)
   )) {
     const multiImagePath = joinFunc('plots', 'image')
-    const isMulti = path.includes(multiImagePath)
-    const pathLabel = path
+    const pathLabel = path.includes(multiImagePath) ? multiImagePath : path
 
     if (!plotAcc[pathLabel]) {
       plotAcc[pathLabel] = {
diff --git a/webview/src/plots/components/App.test.tsx b/webview/src/plots/components/App.test.tsx
index 0c97dc951d..57a4e7b291 100644
--- a/webview/src/plots/components/App.test.tsx
+++ b/webview/src/plots/components/App.test.tsx
@@ -13,6 +13,7 @@ import {
 } from '@testing-library/react'
 import '@testing-library/jest-dom/extend-expect'
 import comparisonTableFixture from 'dvc/src/test/fixtures/plotsDiff/comparison'
+import comparisonTableMultiImgFixture from 'dvc/src/test/fixtures/plotsDiff/comparison/multi'
 import customPlotsFixture from 'dvc/src/test/fixtures/expShow/base/customPlots'
 import plotsRevisionsFixture from 'dvc/src/test/fixtures/plotsDiff/revisions'
 import templatePlotsFixture from 'dvc/src/test/fixtures/plotsDiff/template/webview'
@@ -275,6 +276,70 @@ describe('App', () => {
     expect(loading).toHaveLength(3)
   })
 
+  it('should render loading section states with multi image plots when given a single revision which has not been fetched', async () => {
+    renderAppWithOptionalData({
+      comparison: {
+        height: DEFAULT_PLOT_HEIGHT,
+        plots: [
+          {
+            path: 'training/plots/images/image',
+            revisions: {
+              ad2b5ec: {
+                id: 'ad2b5ec',
+                imgs: [
+                  {
+                    errors: undefined,
+                    loading: true,
+                    url: undefined
+                  },
+                  {
+                    errors: undefined,
+                    loading: true,
+                    url: undefined
+                  },
+                  {
+                    errors: undefined,
+                    loading: true,
+                    url: undefined
+                  }
+                ]
+              }
+            }
+          }
+        ],
+        revisions: [
+          {
+            description: '[exp-a270a]',
+            displayColor: '#945dd6',
+            fetched: false,
+            id: 'ad2b5ec',
+            label: 'ad2b5ec',
+            summaryColumns: []
+          }
+        ],
+        width: DEFAULT_NB_ITEMS_PER_ROW
+      },
+      custom: null,
+      hasPlots: true,
+      hasUnselectedPlots: false,
+      sectionCollapsed: DEFAULT_SECTION_COLLAPSED,
+      selectedRevisions: [
+        {
+          description: '[exp-a270a]',
+          displayColor: '#945dd6',
+          fetched: false,
+          id: 'ad2b5ec',
+          label: 'ad2b5ec',
+          summaryColumns: []
+        }
+      ],
+      template: null
+    })
+    const loading = await screen.findAllByText('Loading...')
+
+    expect(loading).toHaveLength(3)
+  })
+
   it('should render only get started (buttons: add plots, add experiments, add custom plots) when there are some selected exps, all unselected plots, and no custom plots', async () => {
     renderAppWithOptionalData({
       hasPlots: true,
@@ -1384,6 +1449,23 @@ describe('App', () => {
     })
   })
 
+  it('should send a message with the plot path when a comparison table multi img plot is zoomed', () => {
+    renderAppWithOptionalData({
+      comparison: comparisonTableMultiImgFixture
+    })
+
+    const plotWrapper = screen.getAllByTestId('multi-image-cell')[0]
+    const plot = within(plotWrapper).getByTestId('image-plot-button')
+
+    fireEvent.click(plot)
+
+    expect(mockPostMessage).toHaveBeenCalledWith({
+      payload:
+        comparisonTableMultiImgFixture.plots[3].revisions.workspace.imgs[0].url,
+      type: MessageFromWebviewType.ZOOM_PLOT
+    })
+  })
+
   it('should open a modal with the plot zoomed in when clicking a custom plot', () => {
     renderAppWithOptionalData({
       custom: customPlotsFixture
@@ -1570,6 +1652,62 @@ describe('App', () => {
     expect(multiViewPlot).toHaveStyle('--scale: 2')
   })
 
+  describe('Comparison Multi Image Plots', () => {
+    it('should render cells with sliders', () => {
+      renderAppWithOptionalData({
+        comparison: comparisonTableMultiImgFixture
+      })
+
+      const multiImgPlot = screen.getAllByTestId('multi-image-cell')[0]
+      const slider = within(multiImgPlot).getByRole('slider')
+
+      expect(slider).toBeInTheDocument()
+    })
+    it('should update the cell images when the slider changes', () => {
+      renderAppWithOptionalData({
+        comparison: comparisonTableMultiImgFixture
+      })
+
+      const workspaceImgs =
+        comparisonTableMultiImgFixture.plots[3].revisions.workspace.imgs
+      const mainImgs =
+        comparisonTableMultiImgFixture.plots[3].revisions.main.imgs
+
+      const multiImgPlots = screen.getAllByTestId('multi-image-cell')
+      const slider = within(multiImgPlots[0]).getByRole('slider')
+      const workspaceImgEl = within(multiImgPlots[0]).getByRole('img')
+      const mainImgEl = within(multiImgPlots[1]).getByRole('img')
+
+      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[0].url)
+      expect(mainImgEl).toHaveAttribute('src', mainImgs[0].url)
+
+      fireEvent.change(slider, { target: { value: 3 } })
+
+      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[3].url)
+      expect(mainImgEl).toHaveAttribute('src', mainImgs[3].url)
+    })
+
+    it('should disable the multi img row from drag and drop when hovering over a img slider', () => {
+      renderAppWithOptionalData({
+        comparison: comparisonTableMultiImgFixture
+      })
+
+      const multiImgRow = screen.getAllByTestId('comparison-table-body')[3]
+      const multiImgPlots = screen.getAllByTestId('multi-image-cell')
+      const slider = within(multiImgPlots[0]).getByRole('slider')
+
+      expect(multiImgRow.draggable).toBe(true)
+
+      fireEvent.mouseEnter(slider)
+
+      expect(multiImgRow.draggable).toBe(false)
+
+      fireEvent.mouseLeave(slider)
+
+      expect(multiImgRow.draggable).toBe(true)
+    })
+  })
+
   describe('Virtualization', () => {
     const createCustomPlots = (nbOfPlots: number): CustomPlotsData => {
       const plots = []
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
index 14a3826e9a..f1a692f976 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
@@ -44,7 +44,12 @@ export const ComparisionTableRows: React.FC<ComparisonTableRowsProps> = ({
     }
     const revs = plot.revisions
     return (
-      <tbody key={path} id={path} ref={i === 0 ? firstRowRef : undefined}>
+      <tbody
+        data-testid="comparison-table-body"
+        key={path}
+        id={path}
+        ref={i === 0 ? firstRowRef : undefined}
+      >
         <ComparisonTableRow
           path={path}
           plots={columns.map(column => ({
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
index 0065907b8f..2c2f2f6c6a 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
@@ -1,4 +1,4 @@
-import React, { useCallback, MouseEvent, KeyboardEvent } from 'react'
+import React, { useCallback } from 'react'
 import cx from 'classnames'
 import { useDispatch, useSelector } from 'react-redux'
 import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
@@ -33,14 +33,10 @@ export const ComparisonTableMultiCell: React.FC<{
     dispatch(changeDisabledDragIds([]))
   }, [dispatch])
 
-  const disableClick = useCallback((e: MouseEvent | KeyboardEvent) => {
-    e.stopPropagation()
-  }, [])
-
   let imageContent = (
     <button
       className={cx(styles.imageWrapper, styles.multiImageWrapper)}
-      onClick={() => url && zoomPlot(url)}
+      onClick={() => zoomPlot(url)}
       data-testid="image-plot-button"
     >
       <img
@@ -61,16 +57,12 @@ export const ComparisonTableMultiCell: React.FC<{
   }
 
   return (
-    <div className={styles.multiImageWrapper}>
+    <div data-testid="multi-image-cell" className={styles.multiImageWrapper}>
       {imageContent}
-      {/* The div element has children that allow keyboard interaction */}
-      {/* eslint-disable-next-line jsx-a11y/no-static-element-interactions */}
       <div
         className={styles.multiImageSlider}
         onMouseEnter={addDisabled}
         onMouseLeave={removeDisabled}
-        onClick={disableClick}
-        onKeyDown={disableClick}
       >
         <label htmlFor={`${plot.id}-step`}>Step</label>
         <input

From 7f82d988ee353c63489405c41f97b92702ca8018 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 28 Jul 2023 19:59:33 -0500
Subject: [PATCH 09/23] Add missing story

---
 webview/src/stories/ComparisonTable.stories.tsx | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/webview/src/stories/ComparisonTable.stories.tsx b/webview/src/stories/ComparisonTable.stories.tsx
index 448134bc21..172607e779 100644
--- a/webview/src/stories/ComparisonTable.stories.tsx
+++ b/webview/src/stories/ComparisonTable.stories.tsx
@@ -11,6 +11,7 @@ import {
   PlotsComparisonData
 } from 'dvc/src/plots/webview/contract'
 import comparisonTableFixture from 'dvc/src/test/fixtures/plotsDiff/comparison'
+import comparisonTableMultiFixture from 'dvc/src/test/fixtures/plotsDiff/comparison/multi'
 import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
 import { DISABLE_CHROMATIC_SNAPSHOTS } from './util'
 import { ComparisonTable } from '../plots/components/comparisonTable/ComparisonTable'
@@ -102,6 +103,11 @@ const removeImages = (
   return filteredRevisionData
 }
 
+export const WithMultiImages = Template.bind({})
+WithMultiImages.args = {
+  plots: comparisonTableMultiFixture.plots
+}
+
 export const WithMissingData = Template.bind({})
 WithMissingData.args = {
   plots: comparisonTableFixture.plots.map(({ path, revisions }) => ({

From 11490044743af6cebf55c1c82c5324c90e2f43f4 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Sat, 29 Jul 2023 18:15:00 -0500
Subject: [PATCH 10/23] Add tests for paths

---
 extension/src/plots/paths/model.test.ts       | 39 ++++++++++++++++++-
 .../src/test/fixtures/plotsDiff/index.ts      | 12 ++++++
 .../fixtures/plotsDiff/output/multiImage.ts   |  6 +++
 3 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 extension/src/test/fixtures/plotsDiff/output/multiImage.ts

diff --git a/extension/src/plots/paths/model.test.ts b/extension/src/plots/paths/model.test.ts
index fb83d660a7..cf3755dcd0 100644
--- a/extension/src/plots/paths/model.test.ts
+++ b/extension/src/plots/paths/model.test.ts
@@ -2,11 +2,12 @@ import { join } from 'path'
 import { PathsModel } from './model'
 import { PathType } from './collect'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
+import plotsDiffMultiImgFixture from '../../test/fixtures/plotsDiff/output/multiImage'
 import { buildMockMemento } from '../../test/util'
 import { PlotsType, TemplatePlotGroup } from '../webview/contract'
-import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
+import { EXPERIMENT_WORKSPACE_ID, PlotsOutput } from '../../cli/dvc/contract'
 import { ErrorsModel } from '../errors/model'
-import { REVISIONS } from '../../test/fixtures/plotsDiff'
+import { REVISIONS, getMultiImgPaths } from '../../test/fixtures/plotsDiff'
 
 describe('PathsModel', () => {
   const mockDvcRoot = 'test'
@@ -354,6 +355,40 @@ describe('PathsModel', () => {
     expect(model.getComparisonPaths()).toStrictEqual(newOrder)
   })
 
+  it('should sort the order of the comparison paths when there are "by step" images', () => {
+    const model = new PathsModel(
+      mockDvcRoot,
+      buildMockErrorsModel(),
+      buildMockMemento()
+    )
+
+    const singleImgPaths = [
+      join('plots', 'acc.png'),
+      join('plots', 'heatmap.png'),
+      join('plots', 'loss.png')
+    ]
+
+    const currentOrder = [...singleImgPaths, ...getMultiImgPaths()]
+
+    const mixedMultiImgFixture: PlotsOutput = {
+      data: {
+        ...plotsDiffMultiImgFixture.data,
+        [join('plots', 'image', '3.jpg')]:
+          plotsDiffMultiImgFixture.data[join('plots', 'image', '3.jpg')]
+      }
+    }
+    model.transformAndSet(mixedMultiImgFixture, REVISIONS)
+    model.setSelectedRevisions([EXPERIMENT_WORKSPACE_ID])
+
+    expect(model.getComparisonPaths()).toStrictEqual(currentOrder)
+
+    model.setComparisonPathsOrder([join('plots', 'image'), ...singleImgPaths])
+
+    const newOrder = [...getMultiImgPaths(), ...singleImgPaths]
+
+    expect(model.getComparisonPaths()).toStrictEqual(newOrder)
+  })
+
   it('should return the expected children from the test fixture', () => {
     const model = new PathsModel(
       mockDvcRoot,
diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index 1e363babf9..243021902d 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -364,6 +364,14 @@ const multipleVega = (length: number) => {
   return plots
 }
 
+export const getMultiImgPaths = (joinFunc = join) => {
+  const paths = []
+  for (let i = 0; i < 15; i++) {
+    paths.push(joinFunc('plots', 'image', `${i}.jpg`))
+  }
+  return paths
+}
+
 const getMultiImageData = (
   baseUrl: string,
   joinFunc: (...segments: string[]) => string,
@@ -496,6 +504,10 @@ export const getOutput = (baseUrl: string): PlotsOutput => ({
 
 export const getMinimalOutput = (): PlotsOutput => ({ data: { ...basicVega } })
 
+export const getMultiImgOutput = (baseUrl: string): PlotsOutput => ({
+  data: { ...getImageDataWithMultiImgs(baseUrl) }
+})
+
 export const getMultiSourceOutput = (): PlotsOutput => ({
   ...require('./multiSource').default
 })
diff --git a/extension/src/test/fixtures/plotsDiff/output/multiImage.ts b/extension/src/test/fixtures/plotsDiff/output/multiImage.ts
new file mode 100644
index 0000000000..c5262d685e
--- /dev/null
+++ b/extension/src/test/fixtures/plotsDiff/output/multiImage.ts
@@ -0,0 +1,6 @@
+import { basePlotsUrl } from '../../../util'
+import { getMultiImgOutput } from '..'
+
+const data = getMultiImgOutput(basePlotsUrl)
+
+export default data

From 39f9caa7cf047a9df224c93a6225c960e9b9014f Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Sun, 30 Jul 2023 06:28:01 -0500
Subject: [PATCH 11/23] Add vscode tests

---
 .../plotsDiff/comparison/multiVscode.ts       | 25 +++++++++++++++++++
 .../src/test/fixtures/plotsDiff/index.ts      |  2 +-
 extension/src/test/suite/plots/index.test.ts  | 19 ++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)
 create mode 100644 extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts

diff --git a/extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts b/extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts
new file mode 100644
index 0000000000..0a6144b495
--- /dev/null
+++ b/extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts
@@ -0,0 +1,25 @@
+import { getComparisonWebviewMessage } from '..'
+import { Uri, ViewColumn, window } from 'vscode'
+import { ViewKey } from '../../../../webview/constants'
+import { basePlotsUrl } from '../../../util'
+
+const webviewPanel = window.createWebviewPanel(
+  ViewKey.PLOTS,
+  'webview for asWebviewUri',
+  ViewColumn.Active,
+  {
+    enableScripts: true
+  }
+)
+
+const baseUrl = webviewPanel.webview
+  .asWebviewUri(Uri.file(basePlotsUrl))
+  .toString()
+
+webviewPanel.dispose()
+
+const uriJoin = (...segments: string[]) => segments.join('/')
+
+const data = getComparisonWebviewMessage(baseUrl, uriJoin, true)
+
+export default data
diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index 243021902d..f2669e9ce5 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -385,7 +385,7 @@ const getMultiImageData = (
     }[]
   } = {}
   for (let i = 0; i < 15; i++) {
-    const key = join('plots', 'image', `${i}.jpg`)
+    const key = joinFunc('plots', 'image', `${i}.jpg`)
     const values = []
     for (const revision of revisions) {
       values.push({
diff --git a/extension/src/test/suite/plots/index.test.ts b/extension/src/test/suite/plots/index.test.ts
index 826e7f2fdf..f49c1d3c6b 100644
--- a/extension/src/test/suite/plots/index.test.ts
+++ b/extension/src/test/suite/plots/index.test.ts
@@ -11,9 +11,11 @@ import gitLogFixture from '../../fixtures/expShow/base/gitLog'
 import rowOrderFixture from '../../fixtures/expShow/base/rowOrder'
 import customPlotsFixture from '../../fixtures/expShow/base/customPlots'
 import plotsDiffFixture from '../../fixtures/plotsDiff/output'
+import multiImagePlotsDiffFixture from '../../fixtures/plotsDiff/output/multiImage'
 import multiSourcePlotsDiffFixture from '../../fixtures/plotsDiff/multiSource'
 import templatePlotsFixture from '../../fixtures/plotsDiff/template'
 import comparisonPlotsFixture from '../../fixtures/plotsDiff/comparison/vscode'
+import comparisonPlotsMultiImgFixture from '../../fixtures/plotsDiff/comparison/multiVscode'
 import plotsRevisionsFixture from '../../fixtures/plotsDiff/revisions'
 import {
   bypassProcessManagerDebounce,
@@ -968,6 +970,23 @@ suite('Plots Test Suite', () => {
       }
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
+    it('should send the correct data to the webview for multi image plots', async () => {
+      const { plots, messageSpy, mockPlotsDiff } = await buildPlots({
+        disposer: disposable,
+        plotsDiff: multiImagePlotsDiffFixture
+      })
+
+      const webview = await plots.showWebview()
+      await webview.isReady()
+
+      expect(mockPlotsDiff).to.be.called
+
+      const { comparison: comparisonData } = getFirstArgOfLastCall(messageSpy)
+      // we need to grab our fixture from vscode
+      // and edit the images from there
+      expect(comparisonData).to.deep.equal(comparisonPlotsMultiImgFixture)
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should handle a toggle experiment message from the webview', async () => {
       const { plots, experiments } = await buildPlots({
         disposer: disposable,

From 32ac3fc738b7fde776e7b98945d61aaf960bb60d Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Sun, 30 Jul 2023 14:37:18 -0500
Subject: [PATCH 12/23] Fix typos

---
 extension/src/test/suite/plots/index.test.ts                  | 3 +--
 .../plots/components/comparisonTable/ComparisonTable.test.tsx | 4 ----
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/extension/src/test/suite/plots/index.test.ts b/extension/src/test/suite/plots/index.test.ts
index f49c1d3c6b..aa8387837a 100644
--- a/extension/src/test/suite/plots/index.test.ts
+++ b/extension/src/test/suite/plots/index.test.ts
@@ -982,8 +982,7 @@ suite('Plots Test Suite', () => {
       expect(mockPlotsDiff).to.be.called
 
       const { comparison: comparisonData } = getFirstArgOfLastCall(messageSpy)
-      // we need to grab our fixture from vscode
-      // and edit the images from there
+
       expect(comparisonData).to.deep.equal(comparisonPlotsMultiImgFixture)
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
index 61d874182e..5739ddc2e1 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
@@ -280,9 +280,7 @@ describe('ComparisonTable', () => {
             imgs: [
               {
                 errors: undefined,
-                id: revisionWithNoData,
                 loading: false,
-                path: '',
                 url: undefined
               }
             ]
@@ -319,9 +317,7 @@ describe('ComparisonTable', () => {
             imgs: [
               {
                 errors: ['this is an error'],
-                id: revisionWithNoData,
                 loading: false,
-                path: '',
                 url: undefined
               }
             ]

From a98e068101aa85f2a10eccf5ae2b607eb34faf58 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Sun, 30 Jul 2023 14:44:17 -0500
Subject: [PATCH 13/23] Try to fix vscode tests

---
 extension/src/test/fixtures/plotsDiff/index.ts | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index f2669e9ce5..d1a3b05313 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -821,8 +821,7 @@ export const getComparisonWebviewMessage = (
       ? getImageDataWithMultiImgs(baseUrl, joinFunc)
       : getImageData(baseUrl, joinFunc)
   )) {
-    const multiImagePath = joinFunc('plots', 'image')
-    const pathLabel = path.includes(multiImagePath) ? multiImagePath : path
+    const pathLabel = path.includes('image') ? join('plots', 'image') : path
 
     if (!plotAcc[pathLabel]) {
       plotAcc[pathLabel] = {

From a5ded1726deb3a68cc18f693655874fe9f20000d Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 31 Jul 2023 11:35:31 -0500
Subject: [PATCH 14/23] Save paths in a directory

---
 extension/src/cli/dvc/constants.ts            |  2 +-
 extension/src/plots/model/collect.test.ts     | 35 ++++++++++++++--
 extension/src/plots/model/collect.ts          | 22 ++++++++--
 extension/src/plots/paths/collect.test.ts     | 42 +++++++++++++++++++
 extension/src/plots/paths/collect.ts          | 31 ++++++++++++--
 extension/src/plots/paths/model.test.ts       | 27 ++++--------
 extension/src/plots/paths/model.ts            | 40 ++----------------
 .../src/test/fixtures/plotsDiff/index.ts      |  8 ----
 webview/src/plots/components/App.test.tsx     | 41 ++++++++++++++++++
 .../cell/ComparisonTableMultiCell.tsx         |  9 +++-
 10 files changed, 181 insertions(+), 76 deletions(-)

diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts
index bb8f754959..9d29c6b569 100644
--- a/extension/src/cli/dvc/constants.ts
+++ b/extension/src/cli/dvc/constants.ts
@@ -18,7 +18,7 @@ export const EXP_RWLOCK_FILE = join(TEMP_EXP_DIR, 'rwlock.lock')
 export const DEFAULT_NUM_OF_COMMITS_TO_SHOW = 3
 export const NUM_OF_COMMITS_TO_INCREASE = 2
 
-export const MULTI_IMAGE_PATH_REG = /image[/\\|]\d+\.[a-z]+$/i
+export const MULTI_IMAGE_PATH_REG = /\w+[/\\|]\d+\.[a-z]+$/i
 
 export enum Command {
   ADD = 'add',
diff --git a/extension/src/plots/model/collect.test.ts b/extension/src/plots/model/collect.test.ts
index 2773c1bce1..ef7e7bf97f 100644
--- a/extension/src/plots/model/collect.test.ts
+++ b/extension/src/plots/model/collect.test.ts
@@ -8,6 +8,7 @@ import {
   collectImageUrl
 } from './collect'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
+import multiImagePlotsDiffFixture from '../../test/fixtures/plotsDiff/output/multiImage'
 import customPlotsFixture, {
   customPlotsOrderFixture,
   experimentsWithCommits
@@ -91,9 +92,37 @@ describe('collectData', () => {
 
     expect(testBranchHeatmap).toBeDefined()
     expect(testBranchHeatmap).toStrictEqual([
-      plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
-        sameContents(revisions as string[], ['test-branch'])
-      )
+      {
+        ...plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
+          sameContents(revisions as string[], ['test-branch'])
+        ),
+        path: heatmapPlot
+      }
+    ])
+  })
+
+  it('should return the expected output from the comparison multi image test fixture', () => {
+    const { comparisonData } = collectData(multiImagePlotsDiffFixture)
+
+    const heatmapPlot = join('plots', 'heatmap.png')
+
+    expect(Object.keys(comparisonData.main)).toStrictEqual([
+      join('plots', 'acc.png'),
+      heatmapPlot,
+      join('plots', 'loss.png'),
+      join('plots', 'image')
+    ])
+
+    const testBranchHeatmap = comparisonData['test-branch'][heatmapPlot]
+
+    expect(testBranchHeatmap).toBeDefined()
+    expect(testBranchHeatmap).toStrictEqual([
+      {
+        ...plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
+          sameContents(revisions as string[], ['test-branch'])
+        ),
+        path: heatmapPlot
+      }
     ])
   })
 })
diff --git a/extension/src/plots/model/collect.ts b/extension/src/plots/model/collect.ts
index 50ad22c914..01f3c1508f 100644
--- a/extension/src/plots/model/collect.ts
+++ b/extension/src/plots/model/collect.ts
@@ -36,8 +36,8 @@ import {
 import { StrokeDashEncoding } from '../multiSource/constants'
 import { exists } from '../../fileSystem'
 import { hasKey } from '../../util/object'
-import { getParent, getPathArray } from '../../fileSystem/util'
 import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
+import { getParent, getPathArray } from '../../fileSystem/util'
 
 export const getCustomPlotId = (metric: string, param: string) =>
   `custom-${metric}-${param}`
@@ -132,7 +132,7 @@ export type RevisionData = {
 
 export type ComparisonData = {
   [label: string]: {
-    [path: string]: ImagePlot[]
+    [path: string]: (ImagePlot & { path: string })[]
   }
 }
 
@@ -141,7 +141,9 @@ const collectImageData = (
   path: string,
   plot: ImagePlot
 ) => {
-  const pathLabel = path
+  const pathLabel = MULTI_IMAGE_PATH_REG.test(path)
+    ? (getParent(getPathArray(path), 0) as string)
+    : path
   const id = plot.revisions?.[0]
   if (!id) {
     return
@@ -155,7 +157,7 @@ const collectImageData = (
     acc[id][pathLabel] = []
   }
 
-  acc[id][pathLabel].push(plot)
+  acc[id][pathLabel].push({ ...plot, path })
 }
 
 const collectDatapoints = (
@@ -211,6 +213,16 @@ const collectPathData = (acc: DataAccumulator, path: string, plots: Plot[]) => {
   }
 }
 
+const sortComparisonImgPaths = (acc: DataAccumulator) => {
+  for (const [label, paths] of Object.entries(acc.comparisonData)) {
+    for (const path of Object.keys(paths)) {
+      acc.comparisonData[label][path].sort((img1, img2) =>
+        img1.path.localeCompare(img2.path, undefined, { numeric: true })
+      )
+    }
+  }
+}
+
 export const collectData = (output: PlotsOutput): DataAccumulator => {
   const { data } = output
   const acc = {
@@ -222,6 +234,8 @@ export const collectData = (output: PlotsOutput): DataAccumulator => {
     collectPathData(acc, path, plots)
   }
 
+  sortComparisonImgPaths(acc)
+
   return acc
 }
 
diff --git a/extension/src/plots/paths/collect.test.ts b/extension/src/plots/paths/collect.test.ts
index 21cdfed603..8e3e405181 100644
--- a/extension/src/plots/paths/collect.test.ts
+++ b/extension/src/plots/paths/collect.test.ts
@@ -11,6 +11,7 @@ import {
 } from './collect'
 import { TemplatePlotGroup, PlotsType } from '../webview/contract'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
+import multiImagePlotsDiffFixture from '../../test/fixtures/plotsDiff/output/multiImage'
 import { Shape, StrokeDash } from '../multiSource/constants'
 import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 import { REVISIONS } from '../../test/fixtures/plotsDiff'
@@ -75,6 +76,47 @@ describe('collectPaths', () => {
     ])
   })
 
+  it('should return the expected data from the comparison multi image test fixture', () => {
+    expect(
+      collectPaths([], multiImagePlotsDiffFixture, REVISIONS)
+    ).toStrictEqual([
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'acc.png'),
+        revisions: new Set(REVISIONS),
+        type: new Set(['comparison'])
+      },
+      {
+        hasChildren: true,
+        parentPath: undefined,
+        path: 'plots',
+        revisions: new Set(REVISIONS)
+      },
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'heatmap.png'),
+        revisions: new Set(REVISIONS),
+        type: new Set(['comparison'])
+      },
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'loss.png'),
+        revisions: new Set(REVISIONS),
+        type: new Set(['comparison'])
+      },
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'image'),
+        revisions: new Set(REVISIONS),
+        type: new Set(['comparison'])
+      }
+    ])
+  })
+
   it('should update the revision details when any revision is recollected', () => {
     const [remainingPath] = Object.keys(plotsDiffFixture)
     const collectedPaths = collectPaths([], plotsDiffFixture, REVISIONS)
diff --git a/extension/src/plots/paths/collect.ts b/extension/src/plots/paths/collect.ts
index e75a46bd8a..45c721b71c 100644
--- a/extension/src/plots/paths/collect.ts
+++ b/extension/src/plots/paths/collect.ts
@@ -23,6 +23,7 @@ import {
 } from '../multiSource/constants'
 import { MultiSourceEncoding } from '../multiSource/collect'
 import { truncate } from '../../util/string'
+import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
 
 export enum PathType {
   COMPARISON = 'comparison',
@@ -58,8 +59,13 @@ const collectType = (plots: Plot[]) => {
 const getType = (
   data: PlotsData,
   hasChildren: boolean,
-  path: string
+  path: string,
+  isMultiImgPlot?: boolean
 ): Set<PathType> | undefined => {
+  if (isMultiImgPlot) {
+    return new Set<PathType>([PathType.COMPARISON])
+  }
+
   if (hasChildren) {
     return
   }
@@ -123,7 +129,8 @@ const collectOrderedPath = (
   data: PlotsData,
   revisions: Set<string>,
   pathArray: string[],
-  idx: number
+  idx: number,
+  isMultiImgDir: boolean
 ): PlotPath[] => {
   const path = getPath(pathArray, idx)
 
@@ -147,7 +154,9 @@ const collectOrderedPath = (
     revisions
   }
 
-  const type = getType(data, hasChildren, path)
+  const isMultiImgPlot = isMultiImgDir && idx === pathArray.length
+
+  const type = getType(data, hasChildren, path, isMultiImgPlot)
   if (type) {
     plotPath.type = type
   }
@@ -167,9 +176,23 @@ const addRevisionsToPath = (
   }
 
   const pathArray = getPathArray(path)
+  const isMultiImg =
+    MULTI_IMAGE_PATH_REG.test(path) &&
+    !!getType(data, false, path)?.has(PathType.COMPARISON)
+
+  if (isMultiImg) {
+    pathArray.pop()
+  }
 
   for (let reverseIdx = pathArray.length; reverseIdx > 0; reverseIdx--) {
-    acc = collectOrderedPath(acc, data, revisions, pathArray, reverseIdx)
+    acc = collectOrderedPath(
+      acc,
+      data,
+      revisions,
+      pathArray,
+      reverseIdx,
+      isMultiImg
+    )
   }
   return acc
 }
diff --git a/extension/src/plots/paths/model.test.ts b/extension/src/plots/paths/model.test.ts
index cf3755dcd0..c86d0ee81c 100644
--- a/extension/src/plots/paths/model.test.ts
+++ b/extension/src/plots/paths/model.test.ts
@@ -5,9 +5,9 @@ import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
 import plotsDiffMultiImgFixture from '../../test/fixtures/plotsDiff/output/multiImage'
 import { buildMockMemento } from '../../test/util'
 import { PlotsType, TemplatePlotGroup } from '../webview/contract'
-import { EXPERIMENT_WORKSPACE_ID, PlotsOutput } from '../../cli/dvc/contract'
+import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 import { ErrorsModel } from '../errors/model'
-import { REVISIONS, getMultiImgPaths } from '../../test/fixtures/plotsDiff'
+import { REVISIONS } from '../../test/fixtures/plotsDiff'
 
 describe('PathsModel', () => {
   const mockDvcRoot = 'test'
@@ -355,36 +355,27 @@ describe('PathsModel', () => {
     expect(model.getComparisonPaths()).toStrictEqual(newOrder)
   })
 
-  it('should sort the order of the comparison paths when there are "by step" images', () => {
+  it('should group multi comparison plot path directories', () => {
     const model = new PathsModel(
       mockDvcRoot,
       buildMockErrorsModel(),
       buildMockMemento()
     )
-
-    const singleImgPaths = [
+    const currentOrder = [
       join('plots', 'acc.png'),
       join('plots', 'heatmap.png'),
-      join('plots', 'loss.png')
+      join('plots', 'loss.png'),
+      join('plots', 'image')
     ]
 
-    const currentOrder = [...singleImgPaths, ...getMultiImgPaths()]
-
-    const mixedMultiImgFixture: PlotsOutput = {
-      data: {
-        ...plotsDiffMultiImgFixture.data,
-        [join('plots', 'image', '3.jpg')]:
-          plotsDiffMultiImgFixture.data[join('plots', 'image', '3.jpg')]
-      }
-    }
-    model.transformAndSet(mixedMultiImgFixture, REVISIONS)
+    model.transformAndSet(plotsDiffMultiImgFixture, REVISIONS)
     model.setSelectedRevisions([EXPERIMENT_WORKSPACE_ID])
 
     expect(model.getComparisonPaths()).toStrictEqual(currentOrder)
 
-    model.setComparisonPathsOrder([join('plots', 'image'), ...singleImgPaths])
+    const newOrder = [join('plots', 'image'), ...currentOrder.slice(0, 3)]
 
-    const newOrder = [...getMultiImgPaths(), ...singleImgPaths]
+    model.setComparisonPathsOrder(newOrder)
 
     expect(model.getComparisonPaths()).toStrictEqual(newOrder)
   })
diff --git a/extension/src/plots/paths/model.ts b/extension/src/plots/paths/model.ts
index d5bcb97601..dd8aa3ea7a 100644
--- a/extension/src/plots/paths/model.ts
+++ b/extension/src/plots/paths/model.ts
@@ -18,7 +18,6 @@ import { isDvcError } from '../../cli/dvc/reader'
 import { PlotsOutputOrError } from '../../cli/dvc/contract'
 import { getErrorTooltip } from '../../tree'
 import { ErrorsModel } from '../errors/model'
-import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
 
 export class PathsModel extends PathSelectionModel<PlotPath> {
   private readonly errors: ErrorsModel
@@ -127,21 +126,14 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
   }
 
   public getComparisonPaths() {
-    const { paths } = this.sortComparisonPaths(
+    return performSimpleOrderedUpdate(
+      this.comparisonPathsOrder,
       this.getPathsByType(PathType.COMPARISON)
     )
-
-    return performSimpleOrderedUpdate(this.comparisonPathsOrder, paths)
   }
 
   public setComparisonPathsOrder(order: string[]) {
-    const { multiImagePaths } = this.sortComparisonPaths(
-      this.getPathsByType(PathType.COMPARISON)
-    )
-
-    this.comparisonPathsOrder = order.flatMap(path =>
-      path.endsWith('image') ? multiImagePaths : path
-    )
+    this.comparisonPathsOrder = order
     this.persist(
       PersistenceKey.PLOT_COMPARISON_PATHS_ORDER,
       this.comparisonPathsOrder
@@ -170,32 +162,6 @@ export class PathsModel extends PathSelectionModel<PlotPath> {
       .map(({ path }) => path)
   }
 
-  private sortComparisonPaths(originalPaths: string[]) {
-    const multiImagePaths: string[] = []
-    const paths: string[] = []
-    let multiImageInd
-
-    for (const [ind, path] of originalPaths.entries()) {
-      const isMulti = MULTI_IMAGE_PATH_REG.test(path)
-      if (!isMulti) {
-        paths.push(path)
-        continue
-      }
-      if (!multiImageInd) {
-        multiImageInd = ind
-      }
-      multiImagePaths.push(path)
-    }
-
-    multiImagePaths.sort((path1, path2) =>
-      path1.localeCompare(path2, undefined, { numeric: true })
-    )
-
-    paths.splice(multiImageInd || 0, 0, ...multiImagePaths)
-
-    return { multiImagePaths, paths }
-  }
-
   private filterChildren(path: string | undefined): PlotPath[] {
     return this.data.filter(element => {
       if (!this.hasRevisions(element)) {
diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index d1a3b05313..f441b821bc 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -364,14 +364,6 @@ const multipleVega = (length: number) => {
   return plots
 }
 
-export const getMultiImgPaths = (joinFunc = join) => {
-  const paths = []
-  for (let i = 0; i < 15; i++) {
-    paths.push(joinFunc('plots', 'image', `${i}.jpg`))
-  }
-  return paths
-}
-
 const getMultiImageData = (
   baseUrl: string,
   joinFunc: (...segments: string[]) => string,
diff --git a/webview/src/plots/components/App.test.tsx b/webview/src/plots/components/App.test.tsx
index 57a4e7b291..58ab2db372 100644
--- a/webview/src/plots/components/App.test.tsx
+++ b/webview/src/plots/components/App.test.tsx
@@ -1663,6 +1663,7 @@ describe('App', () => {
 
       expect(slider).toBeInTheDocument()
     })
+
     it('should update the cell images when the slider changes', () => {
       renderAppWithOptionalData({
         comparison: comparisonTableMultiImgFixture
@@ -1687,6 +1688,46 @@ describe('App', () => {
       expect(mainImgEl).toHaveAttribute('src', mainImgs[3].url)
     })
 
+    it('should set the slider on max step if the slider state value is too big', () => {
+      renderAppWithOptionalData({
+        comparison: {
+          ...comparisonTableMultiImgFixture,
+          plots: comparisonTableMultiImgFixture.plots.map(plot =>
+            plot.path.includes('image')
+              ? {
+                  ...plot,
+                  revisions: {
+                    ...plot.revisions,
+                    main: {
+                      ...plot.revisions.main,
+                      imgs: plot.revisions.main.imgs.slice(0, 10)
+                    }
+                  }
+                }
+              : plot
+          )
+        }
+      })
+
+      const workspaceImgs =
+        comparisonTableMultiImgFixture.plots[3].revisions.workspace.imgs
+      const mainImgs =
+        comparisonTableMultiImgFixture.plots[3].revisions.main.imgs
+
+      const multiImgPlots = screen.getAllByTestId('multi-image-cell')
+      const slider = within(multiImgPlots[0]).getByRole('slider')
+      const workspaceImgEl = within(multiImgPlots[0]).getByRole('img')
+      const mainImgEl = within(multiImgPlots[1]).getByRole('img')
+
+      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[0].url)
+      expect(mainImgEl).toHaveAttribute('src', mainImgs[0].url)
+
+      fireEvent.change(slider, { target: { value: 14 } })
+
+      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[14].url)
+      expect(mainImgEl).toHaveAttribute('src', mainImgs[9].url)
+    })
+
     it('should disable the multi img row from drag and drop when hovering over a img slider', () => {
       renderAppWithOptionalData({
         comparison: comparisonTableMultiImgFixture
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
index 2c2f2f6c6a..7b15c8129d 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
@@ -12,6 +12,13 @@ import {
 import { PlotsState } from '../../../store'
 import { zoomPlot } from '../../../util/messages'
 
+const getCurrentStep = (stateStep: number | undefined, imgsLength: number) => {
+  if (!stateStep) {
+    return 0
+  }
+  return stateStep > imgsLength - 1 ? imgsLength - 1 : stateStep
+}
+
 export const ComparisonTableMultiCell: React.FC<{
   path: string
   plot: ComparisonPlot
@@ -20,7 +27,7 @@ export const ComparisonTableMultiCell: React.FC<{
     (state: PlotsState) => state.comparison.multiPlotValues
   )
   const dispatch = useDispatch()
-  const currentStep = multiValues[path] || 0
+  const currentStep = getCurrentStep(multiValues[path], plot.imgs.length)
 
   const { loading, url } = plot.imgs[currentStep]
   const missing = !loading && !url

From 866df90d2377763a9761b88376db10aa2501c21b Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 31 Jul 2023 12:21:20 -0500
Subject: [PATCH 15/23] Fix possible error being thrown

---
 extension/src/plots/model/collect.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/extension/src/plots/model/collect.ts b/extension/src/plots/model/collect.ts
index 01f3c1508f..748c9ebc9b 100644
--- a/extension/src/plots/model/collect.ts
+++ b/extension/src/plots/model/collect.ts
@@ -264,7 +264,7 @@ const collectSelectedPathComparisonPlots = ({
   }
 
   for (const id of selectedRevisionIds) {
-    const imgs = comparisonData[id][path]
+    const imgs = comparisonData[id]?.[path]
     pathRevisions.revisions[id] = {
       id,
       imgs: imgs

From f1723e4990d398a455d56b1e4a2f5edfde13c9dd Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 31 Jul 2023 12:38:51 -0500
Subject: [PATCH 16/23] Fix error styles

---
 webview/src/plots/components/comparisonTable/styles.module.scss | 1 +
 1 file changed, 1 insertion(+)

diff --git a/webview/src/plots/components/comparisonTable/styles.module.scss b/webview/src/plots/components/comparisonTable/styles.module.scss
index 2139e59e63..12ef4ac796 100644
--- a/webview/src/plots/components/comparisonTable/styles.module.scss
+++ b/webview/src/plots/components/comparisonTable/styles.module.scss
@@ -168,6 +168,7 @@ $gap: 4px;
   padding-top: 25%;
   padding-bottom: 25%;
   display: flex;
+  flex-direction: column;
   align-items: center;
   justify-content: center;
 }

From e94143c086531abae96567e382ae6d956f4fb701 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 1 Aug 2023 11:59:56 -0500
Subject: [PATCH 17/23] Refactor improvements

---
 extension/src/fileSystem/util.ts              |  4 +-
 extension/src/plots/model/collect.test.ts     | 18 +++------
 extension/src/plots/model/collect.ts          | 36 ++++++++++++++----
 .../cell/ComparisonTableCell.tsx              |  5 ++-
 .../cell/ComparisonTableMultiCell.tsx         | 37 +++----------------
 5 files changed, 46 insertions(+), 54 deletions(-)

diff --git a/extension/src/fileSystem/util.ts b/extension/src/fileSystem/util.ts
index 2ff2d3b132..384422476d 100644
--- a/extension/src/fileSystem/util.ts
+++ b/extension/src/fileSystem/util.ts
@@ -1,4 +1,4 @@
-import { sep } from 'path'
+import { sep, parse } from 'path'
 
 export const getPathArray = (path: string): string[] => path.split(sep)
 
@@ -18,3 +18,5 @@ export const getParent = (pathArray: string[], idx: number) => {
 
 export const removeTrailingSlash = (path: string): string =>
   path.endsWith(sep) ? path.slice(0, -1) : path
+
+export const getFileNameWithoutExt = (path: string) => parse(path).name
diff --git a/extension/src/plots/model/collect.test.ts b/extension/src/plots/model/collect.test.ts
index ef7e7bf97f..85d8d9650b 100644
--- a/extension/src/plots/model/collect.test.ts
+++ b/extension/src/plots/model/collect.test.ts
@@ -92,12 +92,9 @@ describe('collectData', () => {
 
     expect(testBranchHeatmap).toBeDefined()
     expect(testBranchHeatmap).toStrictEqual([
-      {
-        ...plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
-          sameContents(revisions as string[], ['test-branch'])
-        ),
-        path: heatmapPlot
-      }
+      plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
+        sameContents(revisions as string[], ['test-branch'])
+      )
     ])
   })
 
@@ -117,12 +114,9 @@ describe('collectData', () => {
 
     expect(testBranchHeatmap).toBeDefined()
     expect(testBranchHeatmap).toStrictEqual([
-      {
-        ...plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
-          sameContents(revisions as string[], ['test-branch'])
-        ),
-        path: heatmapPlot
-      }
+      plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
+        sameContents(revisions as string[], ['test-branch'])
+      )
     ])
   })
 })
diff --git a/extension/src/plots/model/collect.ts b/extension/src/plots/model/collect.ts
index 748c9ebc9b..67222bdfd5 100644
--- a/extension/src/plots/model/collect.ts
+++ b/extension/src/plots/model/collect.ts
@@ -37,7 +37,11 @@ import { StrokeDashEncoding } from '../multiSource/constants'
 import { exists } from '../../fileSystem'
 import { hasKey } from '../../util/object'
 import { MULTI_IMAGE_PATH_REG } from '../../cli/dvc/constants'
-import { getParent, getPathArray } from '../../fileSystem/util'
+import {
+  getFileNameWithoutExt,
+  getParent,
+  getPathArray
+} from '../../fileSystem/util'
 
 export const getCustomPlotId = (metric: string, param: string) =>
   `custom-${metric}-${param}`
@@ -130,21 +134,31 @@ export type RevisionData = {
   [label: string]: RevisionPathData
 }
 
+type ComparisonDataImgPlot = ImagePlot & { ind?: number }
+
 export type ComparisonData = {
   [label: string]: {
-    [path: string]: (ImagePlot & { path: string })[]
+    [path: string]: ComparisonDataImgPlot[]
   }
 }
 
+const getMultiImagePath = (path: string) =>
+  getParent(getPathArray(path), 0) as string
+
+const getMultiImageInd = (path: string) => {
+  const fileName = getFileNameWithoutExt(path)
+  return Number(fileName)
+}
+
 const collectImageData = (
   acc: ComparisonData,
   path: string,
   plot: ImagePlot
 ) => {
-  const pathLabel = MULTI_IMAGE_PATH_REG.test(path)
-    ? (getParent(getPathArray(path), 0) as string)
-    : path
+  const isMultiImgPlot = MULTI_IMAGE_PATH_REG.test(path)
+  const pathLabel = isMultiImgPlot ? getMultiImagePath(path) : path
   const id = plot.revisions?.[0]
+
   if (!id) {
     return
   }
@@ -157,7 +171,13 @@ const collectImageData = (
     acc[id][pathLabel] = []
   }
 
-  acc[id][pathLabel].push({ ...plot, path })
+  const imgPlot: ComparisonDataImgPlot = { ...plot }
+
+  if (isMultiImgPlot) {
+    imgPlot.ind = getMultiImageInd(path)
+  }
+
+  acc[id][pathLabel].push(imgPlot)
 }
 
 const collectDatapoints = (
@@ -216,8 +236,8 @@ const collectPathData = (acc: DataAccumulator, path: string, plots: Plot[]) => {
 const sortComparisonImgPaths = (acc: DataAccumulator) => {
   for (const [label, paths] of Object.entries(acc.comparisonData)) {
     for (const path of Object.keys(paths)) {
-      acc.comparisonData[label][path].sort((img1, img2) =>
-        img1.path.localeCompare(img2.path, undefined, { numeric: true })
+      acc.comparisonData[label][path].sort(
+        (img1, img2) => (img1.ind || 0) - (img2.ind || 0)
       )
     }
   }
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
index 63158f62fc..53bee7782e 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableCell.tsx
@@ -8,7 +8,8 @@ import { zoomPlot } from '../../../util/messages'
 export const ComparisonTableCell: React.FC<{
   path: string
   plot: ComparisonPlot
-}> = ({ path, plot }) => {
+  imgAlt?: string
+}> = ({ path, plot, imgAlt }) => {
   const plotImg = plot.imgs[0]
 
   const loading = plotImg.loading
@@ -32,7 +33,7 @@ export const ComparisonTableCell: React.FC<{
         className={styles.image}
         draggable={false}
         src={plotImg.url}
-        alt={`Plot of ${path} (${plot.id})`}
+        alt={imgAlt || `Plot of ${path} (${plot.id})`}
       />
     </button>
   )
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
index 7b15c8129d..bc203fca05 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
@@ -1,16 +1,13 @@
 import React, { useCallback } from 'react'
-import cx from 'classnames'
 import { useDispatch, useSelector } from 'react-redux'
 import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
-import { ComparisonTableLoadingCell } from './ComparisonTableLoadingCell'
-import { ComparisonTableMissingCell } from './ComparisonTableMissingCell'
+import { ComparisonTableCell } from './ComparisonTableCell'
 import styles from '../styles.module.scss'
 import {
   changeDisabledDragIds,
   setMultiPlotValue
 } from '../comparisonTableSlice'
 import { PlotsState } from '../../../store'
-import { zoomPlot } from '../../../util/messages'
 
 const getCurrentStep = (stateStep: number | undefined, imgsLength: number) => {
   if (!stateStep) {
@@ -29,9 +26,6 @@ export const ComparisonTableMultiCell: React.FC<{
   const dispatch = useDispatch()
   const currentStep = getCurrentStep(multiValues[path], plot.imgs.length)
 
-  const { loading, url } = plot.imgs[currentStep]
-  const missing = !loading && !url
-
   const addDisabled = useCallback(() => {
     dispatch(changeDisabledDragIds([path]))
   }, [dispatch, path])
@@ -40,32 +34,13 @@ export const ComparisonTableMultiCell: React.FC<{
     dispatch(changeDisabledDragIds([]))
   }, [dispatch])
 
-  let imageContent = (
-    <button
-      className={cx(styles.imageWrapper, styles.multiImageWrapper)}
-      onClick={() => zoomPlot(url)}
-      data-testid="image-plot-button"
-    >
-      <img
-        className={styles.image}
-        draggable={false}
-        src={url}
-        alt={`${currentStep} of ${path} (${plot.id})`}
-      />
-    </button>
-  )
-
-  if (loading) {
-    imageContent = <ComparisonTableLoadingCell />
-  }
-
-  if (missing) {
-    imageContent = <ComparisonTableMissingCell plot={plot.imgs[currentStep]} />
-  }
-
   return (
     <div data-testid="multi-image-cell" className={styles.multiImageWrapper}>
-      {imageContent}
+      <ComparisonTableCell
+        path={path}
+        plot={{ id: plot.id, imgs: [plot.imgs[currentStep]] }}
+        imgAlt={`${currentStep} of ${path} (${plot.id})`}
+      />
       <div
         className={styles.multiImageSlider}
         onMouseEnter={addDisabled}

From f594ec1cf37d3a21af0831e13199511164cb9779 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 2 Aug 2023 10:31:38 -0500
Subject: [PATCH 18/23] Try to improve collectOrderedPath

---
 extension/src/plots/paths/collect.ts | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/extension/src/plots/paths/collect.ts b/extension/src/plots/paths/collect.ts
index 45c721b71c..6c3c6349e5 100644
--- a/extension/src/plots/paths/collect.ts
+++ b/extension/src/plots/paths/collect.ts
@@ -154,7 +154,8 @@ const collectOrderedPath = (
     revisions
   }
 
-  const isMultiImgPlot = isMultiImgDir && idx === pathArray.length
+  const isPathLeaf = idx === pathArray.length
+  const isMultiImgPlot = isMultiImgDir && isPathLeaf
 
   const type = getType(data, hasChildren, path, isMultiImgPlot)
   if (type) {

From e8686300f594f8479cbc440f2bcd9032e153606f Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 2 Aug 2023 10:42:42 -0500
Subject: [PATCH 19/23] Remove slider sync

---
 webview/src/plots/components/App.test.tsx     | 47 +------------------
 .../cell/ComparisonTableMultiCell.tsx         | 26 ++--------
 .../comparisonTable/comparisonTableSlice.ts   |  9 ----
 3 files changed, 6 insertions(+), 76 deletions(-)

diff --git a/webview/src/plots/components/App.test.tsx b/webview/src/plots/components/App.test.tsx
index 58ab2db372..252e64763c 100644
--- a/webview/src/plots/components/App.test.tsx
+++ b/webview/src/plots/components/App.test.tsx
@@ -1664,68 +1664,23 @@ describe('App', () => {
       expect(slider).toBeInTheDocument()
     })
 
-    it('should update the cell images when the slider changes', () => {
+    it('should update the cell image when the slider changes', () => {
       renderAppWithOptionalData({
         comparison: comparisonTableMultiImgFixture
       })
 
       const workspaceImgs =
         comparisonTableMultiImgFixture.plots[3].revisions.workspace.imgs
-      const mainImgs =
-        comparisonTableMultiImgFixture.plots[3].revisions.main.imgs
 
       const multiImgPlots = screen.getAllByTestId('multi-image-cell')
       const slider = within(multiImgPlots[0]).getByRole('slider')
       const workspaceImgEl = within(multiImgPlots[0]).getByRole('img')
-      const mainImgEl = within(multiImgPlots[1]).getByRole('img')
 
       expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[0].url)
-      expect(mainImgEl).toHaveAttribute('src', mainImgs[0].url)
 
       fireEvent.change(slider, { target: { value: 3 } })
 
       expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[3].url)
-      expect(mainImgEl).toHaveAttribute('src', mainImgs[3].url)
-    })
-
-    it('should set the slider on max step if the slider state value is too big', () => {
-      renderAppWithOptionalData({
-        comparison: {
-          ...comparisonTableMultiImgFixture,
-          plots: comparisonTableMultiImgFixture.plots.map(plot =>
-            plot.path.includes('image')
-              ? {
-                  ...plot,
-                  revisions: {
-                    ...plot.revisions,
-                    main: {
-                      ...plot.revisions.main,
-                      imgs: plot.revisions.main.imgs.slice(0, 10)
-                    }
-                  }
-                }
-              : plot
-          )
-        }
-      })
-
-      const workspaceImgs =
-        comparisonTableMultiImgFixture.plots[3].revisions.workspace.imgs
-      const mainImgs =
-        comparisonTableMultiImgFixture.plots[3].revisions.main.imgs
-
-      const multiImgPlots = screen.getAllByTestId('multi-image-cell')
-      const slider = within(multiImgPlots[0]).getByRole('slider')
-      const workspaceImgEl = within(multiImgPlots[0]).getByRole('img')
-      const mainImgEl = within(multiImgPlots[1]).getByRole('img')
-
-      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[0].url)
-      expect(mainImgEl).toHaveAttribute('src', mainImgs[0].url)
-
-      fireEvent.change(slider, { target: { value: 14 } })
-
-      expect(workspaceImgEl).toHaveAttribute('src', workspaceImgs[14].url)
-      expect(mainImgEl).toHaveAttribute('src', mainImgs[9].url)
     })
 
     it('should disable the multi img row from drag and drop when hovering over a img slider', () => {
diff --git a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
index bc203fca05..5700b39505 100644
--- a/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
+++ b/webview/src/plots/components/comparisonTable/cell/ComparisonTableMultiCell.tsx
@@ -1,30 +1,16 @@
-import React, { useCallback } from 'react'
-import { useDispatch, useSelector } from 'react-redux'
+import React, { useCallback, useState } from 'react'
+import { useDispatch } from 'react-redux'
 import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
 import { ComparisonTableCell } from './ComparisonTableCell'
 import styles from '../styles.module.scss'
-import {
-  changeDisabledDragIds,
-  setMultiPlotValue
-} from '../comparisonTableSlice'
-import { PlotsState } from '../../../store'
-
-const getCurrentStep = (stateStep: number | undefined, imgsLength: number) => {
-  if (!stateStep) {
-    return 0
-  }
-  return stateStep > imgsLength - 1 ? imgsLength - 1 : stateStep
-}
+import { changeDisabledDragIds } from '../comparisonTableSlice'
 
 export const ComparisonTableMultiCell: React.FC<{
   path: string
   plot: ComparisonPlot
 }> = ({ path, plot }) => {
-  const multiValues = useSelector(
-    (state: PlotsState) => state.comparison.multiPlotValues
-  )
+  const [currentStep, setCurrentStep] = useState<number>(0)
   const dispatch = useDispatch()
-  const currentStep = getCurrentStep(multiValues[path], plot.imgs.length)
 
   const addDisabled = useCallback(() => {
     dispatch(changeDisabledDragIds([path]))
@@ -54,9 +40,7 @@ export const ComparisonTableMultiCell: React.FC<{
           value={currentStep}
           type="range"
           onChange={event => {
-            dispatch(
-              setMultiPlotValue({ path, value: Number(event.target.value) })
-            )
+            setCurrentStep(Number(event.target.value))
           }}
         />
         <p>{currentStep}</p>
diff --git a/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts b/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts
index 930e85ba59..81029d8710 100644
--- a/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts
+++ b/webview/src/plots/components/comparisonTable/comparisonTableSlice.ts
@@ -11,7 +11,6 @@ export interface ComparisonTableState extends PlotsComparisonData {
   isCollapsed: boolean
   hasData: boolean
   rowHeight: number
-  multiPlotValues: { [path: string]: number }
   disabledDragPlotIds: string[]
 }
 
@@ -22,7 +21,6 @@ export const comparisonTableInitialState: ComparisonTableState = {
   hasData: false,
   height: DEFAULT_HEIGHT[PlotsSection.COMPARISON_TABLE],
   isCollapsed: DEFAULT_SECTION_COLLAPSED[PlotsSection.COMPARISON_TABLE],
-  multiPlotValues: {},
   plots: [],
   revisions: [],
   rowHeight: DEFAULT_ROW_HEIGHT,
@@ -49,12 +47,6 @@ export const comparisonTableSlice = createSlice({
     setCollapsed: (state, action: PayloadAction<boolean>) => {
       state.isCollapsed = action.payload
     },
-    setMultiPlotValue: (
-      state,
-      action: PayloadAction<{ path: string; value: number }>
-    ) => {
-      state.multiPlotValues[action.payload.path] = action.payload.value
-    },
     update: (state, action: PayloadAction<PlotsComparisonData>) => {
       if (!action.payload) {
         return comparisonTableInitialState
@@ -71,7 +63,6 @@ export const comparisonTableSlice = createSlice({
 export const {
   update,
   setCollapsed,
-  setMultiPlotValue,
   changeSize,
   changeDisabledDragIds,
   changeRowHeight

From 6be47b41fa9796e9ceeb24d1efa83f1c8c6bcac5 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 2 Aug 2023 11:12:22 -0500
Subject: [PATCH 20/23] Fix typo in test

---
 .../plots/components/comparisonTable/ComparisonTable.test.tsx | 4 ++--
 .../plots/components/comparisonTable/ComparisonTableRows.tsx  | 1 -
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
index 5739ddc2e1..f037e564ed 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
@@ -293,8 +293,8 @@ describe('ComparisonTable', () => {
           description: undefined,
           displayColor: '#f56565',
           fetched: true,
-          id: 'noData',
-          label: revisionWithNoData,
+          id: revisionWithNoData,
+          label: 'noData',
           summaryColumns: []
         }
       ]
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
index f1a692f976..93960b9e3b 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
@@ -53,7 +53,6 @@ export const ComparisionTableRows: React.FC<ComparisonTableRowsProps> = ({
         <ComparisonTableRow
           path={path}
           plots={columns.map(column => ({
-            ...revs[column.id],
             id: column.id,
             imgs: revs[column.id]?.imgs || [
               { errors: undefined, loading: false, url: undefined }

From 7c97263f3d35ee2932aec118248f382e64bc11f1 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 2 Aug 2023 11:51:47 -0500
Subject: [PATCH 21/23] Improve regex and added tests

---
 extension/src/cli/dvc/constants.ts  |  2 +-
 extension/src/cli/dvc/index.test.ts | 50 ++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts
index 9d29c6b569..d913cb44e2 100644
--- a/extension/src/cli/dvc/constants.ts
+++ b/extension/src/cli/dvc/constants.ts
@@ -18,7 +18,7 @@ export const EXP_RWLOCK_FILE = join(TEMP_EXP_DIR, 'rwlock.lock')
 export const DEFAULT_NUM_OF_COMMITS_TO_SHOW = 3
 export const NUM_OF_COMMITS_TO_INCREASE = 2
 
-export const MULTI_IMAGE_PATH_REG = /\w+[/\\|]\d+\.[a-z]+$/i
+export const MULTI_IMAGE_PATH_REG = /[^/]+[/\\]\d+\.[a-z]+$/i
 
 export enum Command {
   ADD = 'add',
diff --git a/extension/src/cli/dvc/index.test.ts b/extension/src/cli/dvc/index.test.ts
index 55272fc296..db818fa305 100644
--- a/extension/src/cli/dvc/index.test.ts
+++ b/extension/src/cli/dvc/index.test.ts
@@ -1,7 +1,8 @@
+import { join } from 'path'
 import { EventEmitter } from 'vscode'
 import { Disposable, Disposer } from '@hediet/std/disposable'
 import { DvcCli } from '.'
-import { Command } from './constants'
+import { Command, MULTI_IMAGE_PATH_REG } from './constants'
 import { CliResult, CliStarted, typeCheckCommands } from '..'
 import { getProcessEnv } from '../../env'
 import { createProcess } from '../../process/execution'
@@ -52,6 +53,53 @@ describe('typeCheckCommands', () => {
   })
 })
 
+describe('Comparison Multi Image Regex', () => {
+  it('should match a nested image group directory', () => {
+    expect(
+      MULTI_IMAGE_PATH_REG.test(
+        join(
+          'extremely',
+          'super',
+          'super',
+          'super',
+          'nested',
+          'image',
+          '768.svg'
+        )
+      )
+    ).toBe(true)
+  })
+
+  it('should match directorys with spaces or special characters', () => {
+    expect(MULTI_IMAGE_PATH_REG.test(join('mis classified', '5.png'))).toBe(
+      true
+    )
+
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified#^', '5.png'))).toBe(
+      true
+    )
+  })
+
+  it('should match different types of images', () => {
+    const imageFormats = ['svg', 'png', 'jpg', 'jpeg']
+    for (const format of imageFormats) {
+      expect(
+        MULTI_IMAGE_PATH_REG.test(join('misclassified', `5.${format}`))
+      ).toBe(true)
+    }
+  })
+
+  it('should not match files that include none digits or do not have an extension', () => {
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified', 'five.png'))).toBe(
+      false
+    )
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified', '5 4.png'))).toBe(
+      false
+    )
+    expect(MULTI_IMAGE_PATH_REG.test(join('misclassified', '5'))).toBe(false)
+  })
+})
+
 describe('executeDvcProcess', () => {
   it('should pass the correct details to the underlying process given no path to the cli or python binary path', async () => {
     const existingPath = joinEnvPath(

From 530f8c284aa167172f5738ac9b33f42cb7d418fa Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 2 Aug 2023 17:17:00 -0500
Subject: [PATCH 22/23] Use multi images in base fixtures

---
 extension/src/plots/model/collect.test.ts     | 22 ---------
 extension/src/plots/paths/collect.test.ts     | 49 +++----------------
 extension/src/plots/paths/model.test.ts       | 38 +++++---------
 .../fixtures/plotsDiff/comparison/multi.ts    |  5 --
 .../plotsDiff/comparison/multiVscode.ts       | 25 ----------
 .../src/test/fixtures/plotsDiff/index.ts      | 19 ++-----
 .../fixtures/plotsDiff/output/multiImage.ts   |  6 ---
 extension/src/test/suite/plots/index.test.ts  | 21 +-------
 webview/src/plots/components/App.test.tsx     | 14 +++---
 .../src/stories/ComparisonTable.stories.tsx   |  6 ---
 10 files changed, 29 insertions(+), 176 deletions(-)
 delete mode 100644 extension/src/test/fixtures/plotsDiff/comparison/multi.ts
 delete mode 100644 extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts
 delete mode 100644 extension/src/test/fixtures/plotsDiff/output/multiImage.ts

diff --git a/extension/src/plots/model/collect.test.ts b/extension/src/plots/model/collect.test.ts
index 85d8d9650b..0ac8afdb8b 100644
--- a/extension/src/plots/model/collect.test.ts
+++ b/extension/src/plots/model/collect.test.ts
@@ -8,7 +8,6 @@ import {
   collectImageUrl
 } from './collect'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
-import multiImagePlotsDiffFixture from '../../test/fixtures/plotsDiff/output/multiImage'
 import customPlotsFixture, {
   customPlotsOrderFixture,
   experimentsWithCommits
@@ -82,27 +81,6 @@ describe('collectData', () => {
 
     const heatmapPlot = join('plots', 'heatmap.png')
 
-    expect(Object.keys(comparisonData.main)).toStrictEqual([
-      join('plots', 'acc.png'),
-      heatmapPlot,
-      join('plots', 'loss.png')
-    ])
-
-    const testBranchHeatmap = comparisonData['test-branch'][heatmapPlot]
-
-    expect(testBranchHeatmap).toBeDefined()
-    expect(testBranchHeatmap).toStrictEqual([
-      plotsDiffFixture.data[heatmapPlot].find(({ revisions }) =>
-        sameContents(revisions as string[], ['test-branch'])
-      )
-    ])
-  })
-
-  it('should return the expected output from the comparison multi image test fixture', () => {
-    const { comparisonData } = collectData(multiImagePlotsDiffFixture)
-
-    const heatmapPlot = join('plots', 'heatmap.png')
-
     expect(Object.keys(comparisonData.main)).toStrictEqual([
       join('plots', 'acc.png'),
       heatmapPlot,
diff --git a/extension/src/plots/paths/collect.test.ts b/extension/src/plots/paths/collect.test.ts
index 8e3e405181..d90320d68d 100644
--- a/extension/src/plots/paths/collect.test.ts
+++ b/extension/src/plots/paths/collect.test.ts
@@ -11,7 +11,6 @@ import {
 } from './collect'
 import { TemplatePlotGroup, PlotsType } from '../webview/contract'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
-import multiImagePlotsDiffFixture from '../../test/fixtures/plotsDiff/output/multiImage'
 import { Shape, StrokeDash } from '../multiSource/constants'
 import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
 import { REVISIONS } from '../../test/fixtures/plotsDiff'
@@ -46,6 +45,13 @@ describe('collectPaths', () => {
         revisions: new Set(REVISIONS),
         type: new Set(['comparison'])
       },
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'image'),
+        revisions: new Set(REVISIONS),
+        type: new Set(['comparison'])
+      },
       {
         hasChildren: false,
         parentPath: 'logs',
@@ -76,47 +82,6 @@ describe('collectPaths', () => {
     ])
   })
 
-  it('should return the expected data from the comparison multi image test fixture', () => {
-    expect(
-      collectPaths([], multiImagePlotsDiffFixture, REVISIONS)
-    ).toStrictEqual([
-      {
-        hasChildren: false,
-        parentPath: 'plots',
-        path: join('plots', 'acc.png'),
-        revisions: new Set(REVISIONS),
-        type: new Set(['comparison'])
-      },
-      {
-        hasChildren: true,
-        parentPath: undefined,
-        path: 'plots',
-        revisions: new Set(REVISIONS)
-      },
-      {
-        hasChildren: false,
-        parentPath: 'plots',
-        path: join('plots', 'heatmap.png'),
-        revisions: new Set(REVISIONS),
-        type: new Set(['comparison'])
-      },
-      {
-        hasChildren: false,
-        parentPath: 'plots',
-        path: join('plots', 'loss.png'),
-        revisions: new Set(REVISIONS),
-        type: new Set(['comparison'])
-      },
-      {
-        hasChildren: false,
-        parentPath: 'plots',
-        path: join('plots', 'image'),
-        revisions: new Set(REVISIONS),
-        type: new Set(['comparison'])
-      }
-    ])
-  })
-
   it('should update the revision details when any revision is recollected', () => {
     const [remainingPath] = Object.keys(plotsDiffFixture)
     const collectedPaths = collectPaths([], plotsDiffFixture, REVISIONS)
diff --git a/extension/src/plots/paths/model.test.ts b/extension/src/plots/paths/model.test.ts
index c86d0ee81c..caf5468886 100644
--- a/extension/src/plots/paths/model.test.ts
+++ b/extension/src/plots/paths/model.test.ts
@@ -2,7 +2,6 @@ import { join } from 'path'
 import { PathsModel } from './model'
 import { PathType } from './collect'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
-import plotsDiffMultiImgFixture from '../../test/fixtures/plotsDiff/output/multiImage'
 import { buildMockMemento } from '../../test/util'
 import { PlotsType, TemplatePlotGroup } from '../webview/contract'
 import { EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
@@ -60,6 +59,14 @@ describe('PathsModel', () => {
         selected: true,
         type: comparisonType
       },
+      {
+        hasChildren: false,
+        parentPath: 'plots',
+        path: join('plots', 'image'),
+        revisions: new Set(REVISIONS),
+        selected: true,
+        type: comparisonType
+      },
       {
         hasChildren: false,
         parentPath: 'logs',
@@ -341,40 +348,17 @@ describe('PathsModel', () => {
     expect(model.getComparisonPaths()).toStrictEqual([
       join('plots', 'acc.png'),
       join('plots', 'heatmap.png'),
-      join('plots', 'loss.png')
+      join('plots', 'loss.png'),
+      join('plots', 'image')
     ])
 
     const newOrder = [
       join('plots', 'heatmap.png'),
       join('plots', 'acc.png'),
-      join('plots', 'loss.png')
-    ]
-
-    model.setComparisonPathsOrder(newOrder)
-
-    expect(model.getComparisonPaths()).toStrictEqual(newOrder)
-  })
-
-  it('should group multi comparison plot path directories', () => {
-    const model = new PathsModel(
-      mockDvcRoot,
-      buildMockErrorsModel(),
-      buildMockMemento()
-    )
-    const currentOrder = [
-      join('plots', 'acc.png'),
-      join('plots', 'heatmap.png'),
       join('plots', 'loss.png'),
       join('plots', 'image')
     ]
 
-    model.transformAndSet(plotsDiffMultiImgFixture, REVISIONS)
-    model.setSelectedRevisions([EXPERIMENT_WORKSPACE_ID])
-
-    expect(model.getComparisonPaths()).toStrictEqual(currentOrder)
-
-    const newOrder = [join('plots', 'image'), ...currentOrder.slice(0, 3)]
-
     model.setComparisonPathsOrder(newOrder)
 
     expect(model.getComparisonPaths()).toStrictEqual(newOrder)
@@ -406,7 +390,7 @@ describe('PathsModel', () => {
         tooltip: undefined
       },
       {
-        descendantStatuses: [2, 2, 2],
+        descendantStatuses: [2, 2, 2, 2],
         hasChildren: true,
         parentPath: undefined,
         path: 'plots',
diff --git a/extension/src/test/fixtures/plotsDiff/comparison/multi.ts b/extension/src/test/fixtures/plotsDiff/comparison/multi.ts
deleted file mode 100644
index 8c8a2bd1ab..0000000000
--- a/extension/src/test/fixtures/plotsDiff/comparison/multi.ts
+++ /dev/null
@@ -1,5 +0,0 @@
-import { getComparisonWebviewMessage } from '..'
-
-const data = getComparisonWebviewMessage('.', undefined, true)
-
-export default data
diff --git a/extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts b/extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts
deleted file mode 100644
index 0a6144b495..0000000000
--- a/extension/src/test/fixtures/plotsDiff/comparison/multiVscode.ts
+++ /dev/null
@@ -1,25 +0,0 @@
-import { getComparisonWebviewMessage } from '..'
-import { Uri, ViewColumn, window } from 'vscode'
-import { ViewKey } from '../../../../webview/constants'
-import { basePlotsUrl } from '../../../util'
-
-const webviewPanel = window.createWebviewPanel(
-  ViewKey.PLOTS,
-  'webview for asWebviewUri',
-  ViewColumn.Active,
-  {
-    enableScripts: true
-  }
-)
-
-const baseUrl = webviewPanel.webview
-  .asWebviewUri(Uri.file(basePlotsUrl))
-  .toString()
-
-webviewPanel.dispose()
-
-const uriJoin = (...segments: string[]) => segments.join('/')
-
-const data = getComparisonWebviewMessage(baseUrl, uriJoin, true)
-
-export default data
diff --git a/extension/src/test/fixtures/plotsDiff/index.ts b/extension/src/test/fixtures/plotsDiff/index.ts
index f441b821bc..48041b3c0e 100644
--- a/extension/src/test/fixtures/plotsDiff/index.ts
+++ b/extension/src/test/fixtures/plotsDiff/index.ts
@@ -472,11 +472,7 @@ const getImageData = (baseUrl: string, joinFunc = join) => ({
       revisions: ['exp-83425'],
       url: joinFunc(baseUrl, '1ba7bcd_plots_loss.png')
     }
-  ]
-})
-
-const getImageDataWithMultiImgs = (baseUrl: string, joinFunc = join) => ({
-  ...getImageData(baseUrl, joinFunc),
+  ],
   ...getMultiImageData(baseUrl, joinFunc, [
     EXPERIMENT_WORKSPACE_ID,
     'main',
@@ -496,10 +492,6 @@ export const getOutput = (baseUrl: string): PlotsOutput => ({
 
 export const getMinimalOutput = (): PlotsOutput => ({ data: { ...basicVega } })
 
-export const getMultiImgOutput = (baseUrl: string): PlotsOutput => ({
-  data: { ...getImageDataWithMultiImgs(baseUrl) }
-})
-
 export const getMultiSourceOutput = (): PlotsOutput => ({
   ...require('./multiSource').default
 })
@@ -801,18 +793,13 @@ export const MOCK_IMAGE_MTIME = 946684800000
 
 export const getComparisonWebviewMessage = (
   baseUrl: string,
-  joinFunc: (...args: string[]) => string = join,
-  addMulti?: boolean
+  joinFunc: (...args: string[]) => string = join
 ): PlotsComparisonData => {
   const plotAcc: {
     [path: string]: { path: string; revisions: ComparisonRevisionData }
   } = {}
 
-  for (const [path, plots] of Object.entries(
-    addMulti
-      ? getImageDataWithMultiImgs(baseUrl, joinFunc)
-      : getImageData(baseUrl, joinFunc)
-  )) {
+  for (const [path, plots] of Object.entries(getImageData(baseUrl, joinFunc))) {
     const pathLabel = path.includes('image') ? join('plots', 'image') : path
 
     if (!plotAcc[pathLabel]) {
diff --git a/extension/src/test/fixtures/plotsDiff/output/multiImage.ts b/extension/src/test/fixtures/plotsDiff/output/multiImage.ts
deleted file mode 100644
index c5262d685e..0000000000
--- a/extension/src/test/fixtures/plotsDiff/output/multiImage.ts
+++ /dev/null
@@ -1,6 +0,0 @@
-import { basePlotsUrl } from '../../../util'
-import { getMultiImgOutput } from '..'
-
-const data = getMultiImgOutput(basePlotsUrl)
-
-export default data
diff --git a/extension/src/test/suite/plots/index.test.ts b/extension/src/test/suite/plots/index.test.ts
index aa8387837a..d4b5c6b135 100644
--- a/extension/src/test/suite/plots/index.test.ts
+++ b/extension/src/test/suite/plots/index.test.ts
@@ -11,11 +11,9 @@ import gitLogFixture from '../../fixtures/expShow/base/gitLog'
 import rowOrderFixture from '../../fixtures/expShow/base/rowOrder'
 import customPlotsFixture from '../../fixtures/expShow/base/customPlots'
 import plotsDiffFixture from '../../fixtures/plotsDiff/output'
-import multiImagePlotsDiffFixture from '../../fixtures/plotsDiff/output/multiImage'
 import multiSourcePlotsDiffFixture from '../../fixtures/plotsDiff/multiSource'
 import templatePlotsFixture from '../../fixtures/plotsDiff/template'
 import comparisonPlotsFixture from '../../fixtures/plotsDiff/comparison/vscode'
-import comparisonPlotsMultiImgFixture from '../../fixtures/plotsDiff/comparison/multiVscode'
 import plotsRevisionsFixture from '../../fixtures/plotsDiff/revisions'
 import {
   bypassProcessManagerDebounce,
@@ -320,7 +318,8 @@ suite('Plots Test Suite', () => {
       const mockComparisonPathsOrder = [
         join('plots', 'acc.png'),
         join('plots', 'heatmap.png'),
-        join('plots', 'loss.png')
+        join('plots', 'loss.png'),
+        join('plots', 'image')
       ]
 
       messageSpy.resetHistory()
@@ -970,22 +969,6 @@ suite('Plots Test Suite', () => {
       }
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
-    it('should send the correct data to the webview for multi image plots', async () => {
-      const { plots, messageSpy, mockPlotsDiff } = await buildPlots({
-        disposer: disposable,
-        plotsDiff: multiImagePlotsDiffFixture
-      })
-
-      const webview = await plots.showWebview()
-      await webview.isReady()
-
-      expect(mockPlotsDiff).to.be.called
-
-      const { comparison: comparisonData } = getFirstArgOfLastCall(messageSpy)
-
-      expect(comparisonData).to.deep.equal(comparisonPlotsMultiImgFixture)
-    }).timeout(WEBVIEW_TEST_TIMEOUT)
-
     it('should handle a toggle experiment message from the webview', async () => {
       const { plots, experiments } = await buildPlots({
         disposer: disposable,
diff --git a/webview/src/plots/components/App.test.tsx b/webview/src/plots/components/App.test.tsx
index 252e64763c..63ae0df46b 100644
--- a/webview/src/plots/components/App.test.tsx
+++ b/webview/src/plots/components/App.test.tsx
@@ -13,7 +13,6 @@ import {
 } from '@testing-library/react'
 import '@testing-library/jest-dom/extend-expect'
 import comparisonTableFixture from 'dvc/src/test/fixtures/plotsDiff/comparison'
-import comparisonTableMultiImgFixture from 'dvc/src/test/fixtures/plotsDiff/comparison/multi'
 import customPlotsFixture from 'dvc/src/test/fixtures/expShow/base/customPlots'
 import plotsRevisionsFixture from 'dvc/src/test/fixtures/plotsDiff/revisions'
 import templatePlotsFixture from 'dvc/src/test/fixtures/plotsDiff/template/webview'
@@ -1451,7 +1450,7 @@ describe('App', () => {
 
   it('should send a message with the plot path when a comparison table multi img plot is zoomed', () => {
     renderAppWithOptionalData({
-      comparison: comparisonTableMultiImgFixture
+      comparison: comparisonTableFixture
     })
 
     const plotWrapper = screen.getAllByTestId('multi-image-cell')[0]
@@ -1460,8 +1459,7 @@ describe('App', () => {
     fireEvent.click(plot)
 
     expect(mockPostMessage).toHaveBeenCalledWith({
-      payload:
-        comparisonTableMultiImgFixture.plots[3].revisions.workspace.imgs[0].url,
+      payload: comparisonTableFixture.plots[3].revisions.workspace.imgs[0].url,
       type: MessageFromWebviewType.ZOOM_PLOT
     })
   })
@@ -1655,7 +1653,7 @@ describe('App', () => {
   describe('Comparison Multi Image Plots', () => {
     it('should render cells with sliders', () => {
       renderAppWithOptionalData({
-        comparison: comparisonTableMultiImgFixture
+        comparison: comparisonTableFixture
       })
 
       const multiImgPlot = screen.getAllByTestId('multi-image-cell')[0]
@@ -1666,11 +1664,11 @@ describe('App', () => {
 
     it('should update the cell image when the slider changes', () => {
       renderAppWithOptionalData({
-        comparison: comparisonTableMultiImgFixture
+        comparison: comparisonTableFixture
       })
 
       const workspaceImgs =
-        comparisonTableMultiImgFixture.plots[3].revisions.workspace.imgs
+        comparisonTableFixture.plots[3].revisions.workspace.imgs
 
       const multiImgPlots = screen.getAllByTestId('multi-image-cell')
       const slider = within(multiImgPlots[0]).getByRole('slider')
@@ -1685,7 +1683,7 @@ describe('App', () => {
 
     it('should disable the multi img row from drag and drop when hovering over a img slider', () => {
       renderAppWithOptionalData({
-        comparison: comparisonTableMultiImgFixture
+        comparison: comparisonTableFixture
       })
 
       const multiImgRow = screen.getAllByTestId('comparison-table-body')[3]
diff --git a/webview/src/stories/ComparisonTable.stories.tsx b/webview/src/stories/ComparisonTable.stories.tsx
index 172607e779..448134bc21 100644
--- a/webview/src/stories/ComparisonTable.stories.tsx
+++ b/webview/src/stories/ComparisonTable.stories.tsx
@@ -11,7 +11,6 @@ import {
   PlotsComparisonData
 } from 'dvc/src/plots/webview/contract'
 import comparisonTableFixture from 'dvc/src/test/fixtures/plotsDiff/comparison'
-import comparisonTableMultiFixture from 'dvc/src/test/fixtures/plotsDiff/comparison/multi'
 import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract'
 import { DISABLE_CHROMATIC_SNAPSHOTS } from './util'
 import { ComparisonTable } from '../plots/components/comparisonTable/ComparisonTable'
@@ -103,11 +102,6 @@ const removeImages = (
   return filteredRevisionData
 }
 
-export const WithMultiImages = Template.bind({})
-WithMultiImages.args = {
-  plots: comparisonTableMultiFixture.plots
-}
-
 export const WithMissingData = Template.bind({})
 WithMissingData.args = {
   plots: comparisonTableFixture.plots.map(({ path, revisions }) => ({

From 9fb3b7ef8fa88b7ec90815d7c783fc369dd20189 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 2 Aug 2023 18:28:26 -0500
Subject: [PATCH 23/23] Improve tests

---
 extension/src/cli/dvc/index.test.ts           |  4 +--
 .../src/stories/ComparisonTable.stories.tsx   | 28 ++++++++++++-------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/extension/src/cli/dvc/index.test.ts b/extension/src/cli/dvc/index.test.ts
index db818fa305..d38f95f0d7 100644
--- a/extension/src/cli/dvc/index.test.ts
+++ b/extension/src/cli/dvc/index.test.ts
@@ -70,7 +70,7 @@ describe('Comparison Multi Image Regex', () => {
     ).toBe(true)
   })
 
-  it('should match directorys with spaces or special characters', () => {
+  it('should match directories with spaces or special characters', () => {
     expect(MULTI_IMAGE_PATH_REG.test(join('mis classified', '5.png'))).toBe(
       true
     )
@@ -89,7 +89,7 @@ describe('Comparison Multi Image Regex', () => {
     }
   })
 
-  it('should not match files that include none digits or do not have an extension', () => {
+  it('should not match files that include none digits or do not have a file extension', () => {
     expect(MULTI_IMAGE_PATH_REG.test(join('misclassified', 'five.png'))).toBe(
       false
     )
diff --git a/webview/src/stories/ComparisonTable.stories.tsx b/webview/src/stories/ComparisonTable.stories.tsx
index 448134bc21..cef1732ad2 100644
--- a/webview/src/stories/ComparisonTable.stories.tsx
+++ b/webview/src/stories/ComparisonTable.stories.tsx
@@ -5,6 +5,7 @@ import { userEvent, within } from '@storybook/testing-library'
 import React from 'react'
 import { Provider, useDispatch } from 'react-redux'
 import {
+  ComparisonPlotImg,
   ComparisonRevisionData,
   DEFAULT_NB_ITEMS_PER_ROW,
   DEFAULT_PLOT_HEIGHT,
@@ -82,18 +83,25 @@ const removeImages = (
         ['main', '4fb124a'].includes(id)) ||
       id === EXPERIMENT_WORKSPACE_ID
     ) {
+      const isMulti = revisionsData[id].imgs.length > 1
       filteredRevisionData[id] = {
         id,
-        imgs: [
-          {
-            errors:
-              id === 'main'
-                ? [`FileNotFoundError: ${path} not found.`]
-                : undefined,
-            loading: false,
-            url: undefined
-          }
-        ]
+        imgs: isMulti
+          ? (Array.from({ length: revisionsData[id].imgs.length }).fill({
+              errors: undefined,
+              loading: false,
+              url: undefined
+            }) as ComparisonPlotImg[])
+          : [
+              {
+                errors:
+                  id === 'main'
+                    ? [`FileNotFoundError: ${path} not found.`]
+                    : undefined,
+                loading: false,
+                url: undefined
+              }
+            ]
       }
       continue
     }