From 55a3172d8e030b261daea05372b59294f256703c Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 2 Oct 2023 11:25:31 -0500
Subject: [PATCH 01/35] first iteration

---
 extension/src/fileSystem/index.test.ts        |  78 ++++++-----
 extension/src/fileSystem/index.ts             |  41 +++++-
 extension/src/fileSystem/util.ts              |   2 +
 extension/src/pipeline/quickPick.test.ts      | 125 +++++++++++-------
 extension/src/pipeline/quickPick.ts           |  91 +++++++++----
 .../src/test/suite/pipeline/index.test.ts     |   5 +-
 extension/src/vscode/resourcePicker.ts        |  24 +++-
 extension/src/vscode/title.ts                 |   2 +
 8 files changed, 255 insertions(+), 113 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 25730afe61..0c39f28a9c 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -23,7 +23,7 @@ import {
   getPidFromFile,
   getEntryFromJsonFile,
   addPlotToDvcYamlFile,
-  loadDataFile
+  loadDataFiles
 } from '.'
 import { dvcDemoPath } from '../test/util'
 import { DOT_DVC } from '../cli/dvc/constants'
@@ -63,17 +63,22 @@ beforeEach(() => {
   jest.resetAllMocks()
 })
 
-describe('loadDataFile', () => {
+describe('loadDataFiles', () => {
   it('should load in csv file contents', async () => {
     const mockCsvContent = ['epoch,acc', '10,0.69', '11,0.345'].join('\n')
 
     mockedReadFileSync.mockReturnValueOnce(mockCsvContent)
 
-    const result = await loadDataFile('values.csv')
+    const result = await loadDataFiles(['values.csv'])
 
     expect(result).toStrictEqual([
-      { acc: 0.69, epoch: 10 },
-      { acc: 0.345, epoch: 11 }
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'values.csv'
+      }
     ])
   })
 
@@ -85,11 +90,16 @@ describe('loadDataFile', () => {
 
     mockedReadFileSync.mockReturnValueOnce(mockJsonContent)
 
-    const result = await loadDataFile('values.json')
+    const result = await loadDataFiles(['values.json'])
 
     expect(result).toStrictEqual([
-      { acc: 0.69, epoch: 10 },
-      { acc: 0.345, epoch: 11 }
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'values.json'
+      }
     ])
   })
 
@@ -98,11 +108,16 @@ describe('loadDataFile', () => {
 
     mockedReadFileSync.mockReturnValueOnce(mockTsvContent)
 
-    const result = await loadDataFile('values.tsv')
+    const result = await loadDataFiles(['values.tsv'])
 
     expect(result).toStrictEqual([
-      { acc: 0.69, epoch: 10 },
-      { acc: 0.345, epoch: 11 }
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'values.tsv'
+      }
     ])
   })
 
@@ -115,15 +130,20 @@ describe('loadDataFile', () => {
 
     mockedReadFileSync.mockReturnValueOnce(mockYamlContent)
 
-    const result = await loadDataFile('dvc.yaml')
+    const result = await loadDataFiles(['dvc.yaml'])
 
-    expect(result).toStrictEqual({
-      stages: {
-        train: {
-          cmd: 'python train.py'
-        }
+    expect(result).toStrictEqual([
+      {
+        data: {
+          stages: {
+            train: {
+              cmd: 'python train.py'
+            }
+          }
+        },
+        file: 'dvc.yaml'
       }
-    })
+    ])
   })
 
   it('should catch any errors thrown during file parsing', async () => {
@@ -133,7 +153,7 @@ describe('loadDataFile', () => {
     })
 
     for (const file of dataFiles) {
-      const resultWithErr = await loadDataFile(file)
+      const resultWithErr = await loadDataFiles([file])
 
       expect(resultWithErr).toStrictEqual(undefined)
     }
@@ -541,10 +561,9 @@ describe('addPlotToDvcYamlFile', () => {
     mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
 
     addPlotToDvcYamlFile('/', {
-      dataFile: '/data.json',
       template: 'simple',
-      x: 'epochs',
-      y: 'accuracy'
+      x: { file: '/data.json', key: 'epochs' },
+      y: { file: '/data.json', key: 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -560,10 +579,9 @@ describe('addPlotToDvcYamlFile', () => {
     mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent.join('\n'))
 
     addPlotToDvcYamlFile('/', {
-      dataFile: '/data.json',
       template: 'simple',
-      x: 'epochs',
-      y: 'accuracy'
+      x: { file: '/data.json', key: 'epochs' },
+      y: { file: '/data.json', key: 'accuracy' }
     })
 
     mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
@@ -583,10 +601,9 @@ describe('addPlotToDvcYamlFile', () => {
     mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
 
     addPlotToDvcYamlFile('/', {
-      dataFile: '/data.json',
       template: 'simple',
-      x: 'epochs',
-      y: 'accuracy'
+      x: { file: '/data.json', key: 'epochs' },
+      y: { file: '/data.json', key: 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -620,10 +637,9 @@ describe('addPlotToDvcYamlFile', () => {
     mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
 
     addPlotToDvcYamlFile('/', {
-      dataFile: '/data.json',
       template: 'simple',
-      x: 'epochs',
-      y: 'accuracy'
+      x: { file: '/data.json', key: 'epochs' },
+      y: { file: '/data.json', key: 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index 2ddfee9289..4e32be1f95 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -214,13 +214,30 @@ const loadYamlAsDoc = (
   }
 }
 
+const getPlotYamlObj = (cwd: string, plot: PlotConfigData) => {
+  const { x, y, template } = plot
+  const usesSingleFile = x.file === y.file
+  if (usesSingleFile) {
+    const dataFile = x.file
+    const plotName = relative(cwd, dataFile)
+    return { [plotName]: { template, x: x.key, y: y.key } }
+  }
+
+  const plotName = `${template}_plot`
+  return {
+    [plotName]: {
+      template,
+      x: { [relative(cwd, x.file)]: x.key },
+      y: { [relative(cwd, y.file)]: y.key }
+    }
+  }
+}
+
 const getPlotsYaml = (
   cwd: string,
   plotObj: PlotConfigData,
   indentSearchLines: string[]
 ) => {
-  const { dataFile, ...plot } = plotObj
-  const plotName = relative(cwd, dataFile)
   const indentReg = /^( +)[^ ]/
   const indentLine = indentSearchLines.find(line => indentReg.test(line)) || ''
   const spacesMatches = indentLine.match(indentReg)
@@ -228,7 +245,7 @@ const getPlotsYaml = (
 
   return yaml
     .stringify(
-      { plots: [{ [plotName]: plot }] },
+      { plots: [getPlotYamlObj(cwd, plotObj)] },
       { indent: spaces ? spaces.length : 2 }
     )
     .split('\n')
@@ -315,7 +332,7 @@ const loadTsv = (path: string) => {
   }
 }
 
-export const loadDataFile = (file: string): unknown => {
+const loadDataFile = (file: string): unknown => {
   const ext = getFileExtension(file)
 
   switch (ext) {
@@ -330,6 +347,22 @@ export const loadDataFile = (file: string): unknown => {
   }
 }
 
+export const loadDataFiles = async (
+  files: string[]
+): Promise<{ file: string; data: unknown }[] | undefined> => {
+  const filesData: { file: string; data: unknown }[] = []
+  for (const file of files) {
+    const data = await loadDataFile(file)
+
+    if (!data) {
+      return undefined
+    }
+
+    filesData.push({ data, file })
+  }
+  return filesData
+}
+
 export const writeJson = <
   T extends Record<string, unknown> | Array<Record<string, unknown>>
 >(
diff --git a/extension/src/fileSystem/util.ts b/extension/src/fileSystem/util.ts
index 36ab159925..53effd7ac0 100644
--- a/extension/src/fileSystem/util.ts
+++ b/extension/src/fileSystem/util.ts
@@ -22,4 +22,6 @@ export const getParent = (pathArray: string[], idx: number) => {
 export const removeTrailingSlash = (path: string): string =>
   path.endsWith(sep) ? path.slice(0, -1) : path
 
+export const getFileName = (path: string) => parse(path).base
+
 export const getFileNameWithoutExt = (path: string) => parse(path).name
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index e43baab160..51d9864113 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -1,14 +1,16 @@
+import { QuickPickItemKind } from 'vscode'
 import { pickPlotConfiguration } from './quickPick'
-import { pickFile } from '../vscode/resourcePicker'
-import { quickPickOne } from '../vscode/quickPick'
-import { getFileExtension, loadDataFile } from '../fileSystem'
+import { pickFiles } from '../vscode/resourcePicker'
+import { quickPickOne, quickPickValue } from '../vscode/quickPick'
+import { getFileExtension, loadDataFiles } from '../fileSystem'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
 
-const mockedPickFile = jest.mocked(pickFile)
-const mockedLoadDataFile = jest.mocked(loadDataFile)
+const mockedPickFiles = jest.mocked(pickFiles)
+const mockedLoadDataFiles = jest.mocked(loadDataFiles)
 const mockedGetFileExt = jest.mocked(getFileExtension)
 const mockedQuickPickOne = jest.mocked(quickPickOne)
+const mockedQuickPickValue = jest.mocked(quickPickValue)
 const mockedToast = jest.mocked(Toast)
 const mockedShowError = jest.fn()
 mockedToast.showError = mockedShowError
@@ -47,18 +49,20 @@ const mockValidData = [
 
 describe('pickPlotConfiguration', () => {
   it('should let the user pick from files with accepted data types', async () => {
-    mockedPickFile.mockResolvedValueOnce('file.json')
-    mockedLoadDataFile.mockReturnValueOnce(mockValidData)
+    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
 
     await pickPlotConfiguration()
 
-    expect(mockedPickFile).toHaveBeenCalledWith(Title.SELECT_PLOT_DATA, {
+    expect(mockedPickFiles).toHaveBeenCalledWith(Title.SELECT_PLOT_DATA, {
       'Data Formats': ['json', 'csv', 'tsv', 'yaml']
     })
   })
 
   it('should return early if user does not select a file', async () => {
-    mockedPickFile.mockResolvedValueOnce(undefined)
+    mockedPickFiles.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration()
 
@@ -66,20 +70,20 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should show a toast message if the file fails to parse', async () => {
-    mockedPickFile.mockResolvedValueOnce('file.csv')
-    mockedLoadDataFile.mockReturnValueOnce(undefined)
+    mockedPickFiles.mockResolvedValueOnce(['data.csv'])
+    mockedLoadDataFiles.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration()
 
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'Failed to parse the requested file. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'Failed to parse the requested file(s). Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
   it('should show a toast message if an array of objects (with atleast two keys) are not found within a file', async () => {
-    mockedPickFile.mockResolvedValue('file.yaml')
+    mockedPickFiles.mockResolvedValue(['file.yaml'])
     const invalidValues: unknown[] = [
       'string',
       13,
@@ -115,20 +119,22 @@ describe('pickPlotConfiguration', () => {
     let result
 
     for (const [ind, invalidVal] of invalidValues.entries()) {
-      mockedLoadDataFile.mockReturnValueOnce(invalidVal)
+      mockedLoadDataFiles.mockResolvedValueOnce([
+        { data: invalidVal, file: 'file.yaml' }
+      ])
 
       result = await pickPlotConfiguration()
 
       expect(result).toStrictEqual(undefined)
       expect(mockedShowError).toHaveBeenCalledTimes(1 + ind)
       expect(mockedShowError).toHaveBeenCalledWith(
-        'The requested file does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+        'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
       )
     }
   })
 
   it('should parse fields from valid data files', async () => {
-    mockedPickFile.mockResolvedValue('file.yaml')
+    mockedPickFiles.mockResolvedValue(['file.yaml'])
     const validValues: unknown[] = [
       [{ field1: 1, field2: 2 }],
       [
@@ -145,7 +151,9 @@ describe('pickPlotConfiguration', () => {
     ]
 
     for (const [ind, val] of validValues.entries()) {
-      mockedLoadDataFile.mockReturnValueOnce(val)
+      mockedLoadDataFiles.mockResolvedValueOnce([
+        { data: val, file: 'file.yaml' }
+      ])
 
       await pickPlotConfiguration()
 
@@ -155,12 +163,14 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should let the user pick a template, x field, and y field', async () => {
-    mockedPickFile.mockResolvedValueOnce('file.json')
-    mockedLoadDataFile.mockReturnValueOnce(mockValidData)
-    mockedQuickPickOne
-      .mockResolvedValueOnce('simple')
-      .mockResolvedValueOnce('actual')
-      .mockResolvedValueOnce('prob')
+    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedQuickPickValue
+      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
 
     const result = await pickPlotConfiguration()
 
@@ -179,27 +189,45 @@ describe('pickPlotConfiguration', () => {
       ],
       'Pick a Plot Template'
     )
-    expect(mockedQuickPickOne).toHaveBeenNthCalledWith(
-      2,
-      ['actual', 'prob'],
-      'Pick a Metric for X'
+    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
+      1,
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
+      ],
+      { title: Title.SELECT_PLOT_X_METRIC }
     )
-    expect(mockedQuickPickOne).toHaveBeenNthCalledWith(
-      3,
-      ['prob'],
-      'Pick a Metric for Y'
+    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
+      2,
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
+      ],
+      {
+        title: Title.SELECT_PLOT_Y_METRIC
+      }
     )
     expect(result).toStrictEqual({
-      dataFile: 'file.json',
       template: 'simple',
-      x: 'actual',
-      y: 'prob'
+      x: { file: 'file.json', key: 'actual' },
+      y: { file: 'file.json', key: 'prob' }
     })
   })
 
   it('should return early if the user does not pick a template', async () => {
-    mockedPickFile.mockResolvedValueOnce('file.json')
-    mockedLoadDataFile.mockReturnValueOnce(mockValidData)
+    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
     mockedQuickPickOne.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration()
@@ -210,29 +238,32 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should return early if the user does not pick a x field', async () => {
-    mockedPickFile.mockResolvedValueOnce('file.json')
-    mockedLoadDataFile.mockReturnValueOnce(mockValidData)
-    mockedQuickPickOne
-      .mockResolvedValueOnce('simple')
-      .mockResolvedValueOnce(undefined)
+    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedQuickPickValue.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration()
 
-    expect(mockedQuickPickOne).toHaveBeenCalledTimes(2)
+    expect(mockedQuickPickValue).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
   })
 
   it('should return early if the user does not pick a y field', async () => {
-    mockedPickFile.mockResolvedValueOnce('file.json')
-    mockedLoadDataFile.mockReturnValueOnce(mockValidData)
-    mockedQuickPickOne
-      .mockResolvedValueOnce('simple')
+    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedQuickPickValue
       .mockResolvedValueOnce('actual')
       .mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration()
 
-    expect(mockedQuickPickOne).toHaveBeenCalledTimes(3)
+    expect(mockedQuickPickValue).toHaveBeenCalledTimes(2)
     expect(result).toStrictEqual(undefined)
   })
 })
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index fd03adf667..5fda8e3b9e 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -1,40 +1,54 @@
 import isEqual from 'lodash.isequal'
+import { QuickPickItemKind } from 'vscode'
 import {
   PLOT_TEMPLATES,
   Value,
   ValueTree,
   isValueTree
 } from '../cli/dvc/contract'
-import { loadDataFile } from '../fileSystem'
-import { quickPickOne } from '../vscode/quickPick'
-import { pickFile } from '../vscode/resourcePicker'
+import { loadDataFiles } from '../fileSystem'
+import { quickPickOne, quickPickValue } from '../vscode/quickPick'
+import { pickFiles } from '../vscode/resourcePicker'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
+import { getFileName } from '../fileSystem/util'
 
-const pickDataFile = () => {
-  return pickFile(Title.SELECT_PLOT_DATA, {
+const pickDataFiles = (): Thenable<string[] | undefined> =>
+  pickFiles(Title.SELECT_PLOT_DATA, {
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
-}
 
-const pickTemplateAndFields = async (
-  fields: string[]
-): Promise<{ x: string; y: string; template: string } | undefined> => {
+const pickTemplateAndFields = async (fields: {
+  [file: string]: string[]
+}): Promise<PlotConfigData | undefined> => {
   const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
 
   if (!template) {
     return
   }
 
-  const x = await quickPickOne(fields, 'Pick a Metric for X')
+  const items = []
+
+  for (const [file, keys] of Object.entries(fields)) {
+    items.push(
+      {
+        kind: QuickPickItemKind.Separator,
+        label: getFileName(file),
+        value: undefined
+      },
+      ...keys.map(key => ({ label: key, value: { file, key } }))
+    )
+  }
+
+  const x = await quickPickValue(items, { title: Title.SELECT_PLOT_X_METRIC })
 
   if (!x) {
     return
   }
 
-  const y = await quickPickOne(
-    fields.filter(field => x !== field),
-    'Pick a Metric for Y'
+  const y = await quickPickValue(
+    items.filter(item => item.value?.key !== x.key),
+    { title: Title.SELECT_PLOT_Y_METRIC }
   )
 
   if (!y) {
@@ -45,10 +59,9 @@ const pickTemplateAndFields = async (
 }
 
 export type PlotConfigData = {
-  dataFile: string
+  x: { file: string; key: string }
   template: string
-  x: string
-  y: string
+  y: { file: string; key: string }
 }
 
 type UnknownValue = Value | ValueTree
@@ -79,28 +92,56 @@ const getFieldOptions = (data: UnknownValue): string[] => {
     : []
 }
 
+const getFieldOptionsFromArr = (
+  dataArr: { data: UnknownValue; file: string }[]
+) => {
+  const keys: {
+    [file: string]: string[]
+  } = {}
+  let keysAmount = 0
+
+  for (const { file, data } of dataArr) {
+    const fields = getFieldOptions(data)
+
+    if (fields.length === 0) {
+      continue
+    }
+
+    keysAmount += fields.length
+
+    if (!keys[file]) {
+      keys[file] = []
+    }
+    keys[file].push(...fields)
+  }
+
+  return { keys, keysAmount }
+}
+
 export const pickPlotConfiguration = async (): Promise<
   PlotConfigData | undefined
 > => {
-  const file = await pickDataFile()
+  const files = await pickDataFiles()
 
-  if (!file) {
+  if (!files) {
     return
   }
 
-  const data = await loadDataFile(file)
+  const filesData = await loadDataFiles(files)
 
-  if (!data) {
+  if (!filesData) {
     return Toast.showError(
-      'Failed to parse the requested file. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'Failed to parse the requested file(s). Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   }
 
-  const keys = getFieldOptions(data as UnknownValue)
+  const { keys, keysAmount } = getFieldOptionsFromArr(
+    filesData as { data: UnknownValue; file: string }[]
+  )
 
-  if (keys.length < 2) {
+  if (keysAmount < 2) {
     return Toast.showError(
-      'The requested file does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   }
 
@@ -110,5 +151,5 @@ export const pickPlotConfiguration = async (): Promise<
     return
   }
 
-  return { ...templateAndFields, dataFile: file }
+  return { ...templateAndFields }
 }
diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts
index bde79ed3ff..9ba5a92915 100644
--- a/extension/src/test/suite/pipeline/index.test.ts
+++ b/extension/src/test/suite/pipeline/index.test.ts
@@ -310,10 +310,9 @@ suite('Pipeline Test Suite', () => {
     expect(mockAddPlotToDvcFile).not.to.be.called
 
     mockPickPlotConfiguration.onSecondCall().resolves({
-      dataFile: 'results.json',
       template: 'simple',
-      x: 'step',
-      y: 'acc'
+      x: { file: 'results.json', key: 'step' },
+      y: { file: 'results.json', key: 'acc' }
     })
 
     await pipeline.addDataSeriesPlot()
diff --git a/extension/src/vscode/resourcePicker.ts b/extension/src/vscode/resourcePicker.ts
index 43b6b5aa15..7d47efb25f 100644
--- a/extension/src/vscode/resourcePicker.ts
+++ b/extension/src/vscode/resourcePicker.ts
@@ -25,12 +25,30 @@ export const pickFile = async (
   }
 }
 
-export const pickResources = (title: Title): Thenable<Uri[] | undefined> => {
-  return window.showOpenDialog({
+export const pickResources = (
+  title: Title,
+  filters?: OpenDialogOptions['filters']
+): Thenable<Uri[] | undefined> => {
+  const opts: OpenDialogOptions = {
     canSelectFiles: true,
     canSelectFolders: true,
     canSelectMany: true,
     openLabel: 'Select',
     title
-  })
+  }
+
+  if (!filters) {
+    opts.filters = filters
+  }
+
+  return window.showOpenDialog(opts)
+}
+
+export const pickFiles = async (
+  title: Title,
+  filters?: OpenDialogOptions['filters']
+): Promise<string[] | undefined> => {
+  const files = await pickResources(title, filters)
+
+  return files?.map(({ fsPath }) => fsPath)
 }
diff --git a/extension/src/vscode/title.ts b/extension/src/vscode/title.ts
index d9fbbf486b..1545c25e83 100644
--- a/extension/src/vscode/title.ts
+++ b/extension/src/vscode/title.ts
@@ -40,6 +40,8 @@ export enum Title {
   SELECT_SORT_DIRECTION = 'Select Sort Direction',
   SELECT_SORTS_TO_REMOVE = 'Select Sort(s) to Remove',
   SELECT_TRAINING_SCRIPT = 'Select your Training Script',
+  SELECT_PLOT_X_METRIC = 'Select a Metric for X',
+  SELECT_PLOT_Y_METRIC = 'Select a Metric for Y',
   SELECT_PLOT_DATA = 'Select Data to Plot',
   SELECT_PLOT_TYPE = 'Select Plot Type',
   SETUP_WORKSPACE = 'Setup the Workspace',

From 68953fc6b4654094caf6f2c2be2505da5a98e6e4 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 2 Oct 2023 11:39:33 -0500
Subject: [PATCH 02/35] adjust filter

---
 extension/src/pipeline/quickPick.ts | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 5fda8e3b9e..7b5eb2f814 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -47,7 +47,9 @@ const pickTemplateAndFields = async (fields: {
   }
 
   const y = await quickPickValue(
-    items.filter(item => item.value?.key !== x.key),
+    items.filter(
+      item => item.value?.key !== x.key && item.value?.file !== x.file
+    ),
     { title: Title.SELECT_PLOT_Y_METRIC }
   )
 

From e410b537fb500e336e6a749c58c434a46e0a23e0 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 2 Oct 2023 18:17:52 -0500
Subject: [PATCH 03/35] Add new tests

---
 extension/src/fileSystem/index.test.ts   |  60 ++++++++++++
 extension/src/fileSystem/index.ts        |   4 +-
 extension/src/pipeline/quickPick.test.ts | 120 ++++++++++++++++++++++-
 extension/src/pipeline/quickPick.ts      |   4 +-
 4 files changed, 180 insertions(+), 8 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 0c39f28a9c..ac1ebf0663 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -146,6 +146,33 @@ describe('loadDataFiles', () => {
     ])
   })
 
+  it('should load in the contents of multiple files', async () => {
+    const mockTsvContent = ['epoch\tacc', '10\t0.69', '11\t0.345'].join('\n')
+    const mockCsvContent = ['epoch2,acc2', '10,0.679', '11,0.3'].join('\n')
+
+    mockedReadFileSync.mockReturnValueOnce(mockTsvContent)
+    mockedReadFileSync.mockReturnValueOnce(mockCsvContent)
+
+    const result = await loadDataFiles(['values.tsv', 'values2.csv'])
+
+    expect(result).toStrictEqual([
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'values.tsv'
+      },
+      {
+        data: [
+          { acc2: 0.679, epoch2: 10 },
+          { acc2: 0.3, epoch2: 11 }
+        ],
+        file: 'values2.csv'
+      }
+    ])
+  })
+
   it('should catch any errors thrown during file parsing', async () => {
     const dataFiles = ['values.csv', 'file.json', 'file.tsv', 'dvc.yaml']
     mockedReadFileSync.mockImplementation(() => {
@@ -158,6 +185,39 @@ describe('loadDataFiles', () => {
       expect(resultWithErr).toStrictEqual(undefined)
     }
   })
+
+  it('should catch any errors thrown during the parsing of multiple files', async () => {
+    const dataFiles = ['values.csv', 'file.tsv', 'file.json']
+    const mockCsvContent = ['epoch,acc', '10,0.69', '11,0.345'].join('\n')
+    const mockJsonContent = JSON.stringify([
+      { acc: 0.69, epoch: 10 },
+      { acc: 0.345, epoch: 11 }
+    ])
+    mockedReadFileSync
+      .mockReturnValueOnce(mockCsvContent)
+      .mockImplementationOnce(() => {
+        throw new Error('fake error')
+      })
+      .mockReturnValueOnce(mockJsonContent)
+
+    const resultWithErr = await loadDataFiles(dataFiles)
+    expect(resultWithErr).toStrictEqual([
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'values.csv'
+      },
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'file.json'
+      }
+    ])
+  })
 })
 
 describe('writeJson', () => {
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index 4e32be1f95..f6bf8e50c3 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -355,12 +355,12 @@ export const loadDataFiles = async (
     const data = await loadDataFile(file)
 
     if (!data) {
-      return undefined
+      continue
     }
 
     filesData.push({ data, file })
   }
-  return filesData
+  return filesData.length === 0 ? undefined : filesData
 }
 
 export const writeJson = <
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 51d9864113..900a4c9cc7 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -69,7 +69,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
   })
 
-  it('should show a toast message if the file fails to parse', async () => {
+  it('should show a toast message if files(s) fail to parse', async () => {
     mockedPickFiles.mockResolvedValueOnce(['data.csv'])
     mockedLoadDataFiles.mockResolvedValueOnce(undefined)
 
@@ -82,7 +82,7 @@ describe('pickPlotConfiguration', () => {
     )
   })
 
-  it('should show a toast message if an array of objects (with atleast two keys) are not found within a file', async () => {
+  it('should show a toast message if an array of objects (with atleast two keys total) are not found within a file', async () => {
     mockedPickFiles.mockResolvedValue(['file.yaml'])
     const invalidValues: unknown[] = [
       'string',
@@ -133,7 +133,28 @@ describe('pickPlotConfiguration', () => {
     }
   })
 
-  it('should parse fields from valid data files', async () => {
+  it('should show a toast message if atleast two fields are not found with multiple files', async () => {
+    mockedPickFiles.mockResolvedValueOnce([
+      'file.yaml',
+      'file2.yaml',
+      'file3.yaml'
+    ])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: '', file: 'file.yaml' },
+      { data: { val: [{ field1: 'only one field' }] }, file: 'file2.yaml' },
+      { data: [], file: 'file3.yaml' }
+    ])
+
+    const result = await pickPlotConfiguration()
+
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+    expect(mockedShowError).toHaveBeenCalledWith(
+      'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+    )
+  })
+
+  it('should parse fields from a valid data file', async () => {
     mockedPickFiles.mockResolvedValue(['file.yaml'])
     const validValues: unknown[] = [
       [{ field1: 1, field2: 2 }],
@@ -162,6 +183,38 @@ describe('pickPlotConfiguration', () => {
     }
   })
 
+  it('should parse fields from a mixture of valid and invalid files if atleast 2 fields are found in total', async () => {
+    mockedPickFiles.mockResolvedValueOnce([
+      'file.yaml',
+      'file2.yaml',
+      'file3.yaml'
+    ])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      {
+        data: [
+          { field1: 1, field2: 2, field3: 1, field4: 2 },
+          { field1: 1, field2: 2, field3: 1, field4: 2 }
+        ],
+        file: 'file.yaml'
+      },
+      {
+        data: {
+          field1: [{ field1: 1 }, { field1: 3 }]
+        },
+        file: 'file3.yaml'
+      },
+      {
+        data: undefined,
+        file: 'file3.yaml'
+      }
+    ])
+
+    await pickPlotConfiguration()
+
+    expect(mockedShowError).not.toHaveBeenCalled()
+    expect(mockedQuickPickOne).toHaveBeenCalledTimes(1)
+  })
+
   it('should let the user pick a template, x field, and y field', async () => {
     mockedPickFiles.mockResolvedValueOnce(['file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
@@ -223,6 +276,67 @@ describe('pickPlotConfiguration', () => {
     })
   })
 
+  it('should let the user pick a x field and y field from multiple files', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['file.json', 'file2.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' },
+      { data: mockValidData, file: 'file2.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedQuickPickValue
+      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: 'file2.json', key: 'prob' })
+
+    const result = await pickPlotConfiguration()
+
+    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
+      1,
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file2.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+      ],
+      { title: Title.SELECT_PLOT_X_METRIC }
+    )
+    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
+      2,
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file2.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+      ],
+      {
+        title: Title.SELECT_PLOT_Y_METRIC
+      }
+    )
+    expect(result).toStrictEqual({
+      template: 'simple',
+      x: { file: 'file.json', key: 'actual' },
+      y: { file: 'file2.json', key: 'prob' }
+    })
+  })
+
   it('should return early if the user does not pick a template', async () => {
     mockedPickFiles.mockResolvedValueOnce(['file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 7b5eb2f814..c131003853 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -47,9 +47,7 @@ const pickTemplateAndFields = async (fields: {
   }
 
   const y = await quickPickValue(
-    items.filter(
-      item => item.value?.key !== x.key && item.value?.file !== x.file
-    ),
+    items.filter(item => !isEqual(item.value, x)),
     { title: Title.SELECT_PLOT_Y_METRIC }
   )
 

From 1a86d95786852ac445c31a80fc98ad2c949aa8fc Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 2 Oct 2023 18:23:13 -0500
Subject: [PATCH 04/35] Add missing test for dvc.yaml

---
 extension/src/fileSystem/index.test.ts | 28 ++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index ac1ebf0663..5e4bc8c2d5 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -632,6 +632,34 @@ describe('addPlotToDvcYamlFile', () => {
     )
   })
 
+  it('should add the new plot with fields coming from different fiels', () => {
+    const mockDvcYamlContent = mockStagesLines.join('\n')
+    const mockPlotYamlContent = [
+      '',
+      'plots:',
+      '  - simple_plot:',
+      '      template: simple',
+      '      x:',
+      '        data.json: epochs',
+      '      y:',
+      '        acc.json: accuracy',
+      ''
+    ].join('\n')
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+
+    addPlotToDvcYamlFile('/', {
+      template: 'simple',
+      x: { file: '/data.json', key: 'epochs' },
+      y: { file: '/acc.json', key: 'accuracy' }
+    })
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockDvcYamlContent + mockPlotYamlContent
+    )
+  })
+
   it('should add the new plot if the dvc.yaml file already has plots', () => {
     const mockDvcYamlContent = [...mockPlotsListLines, ...mockStagesLines]
     const mockPlotYamlContent = [...mockNewPlotLines, '']

From f81ef2f109c40527359a338adef1fcfe630f1d18 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 2 Oct 2023 18:33:53 -0500
Subject: [PATCH 05/35] resourcePicker adjustments

---
 extension/src/vscode/resourcePicker.test.ts | 15 ++++++++++++++-
 extension/src/vscode/resourcePicker.ts      | 10 ++--------
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/extension/src/vscode/resourcePicker.test.ts b/extension/src/vscode/resourcePicker.test.ts
index 23256ffc5f..38ae071e40 100644
--- a/extension/src/vscode/resourcePicker.test.ts
+++ b/extension/src/vscode/resourcePicker.test.ts
@@ -1,6 +1,6 @@
 import { resolve } from 'path'
 import { Uri, window } from 'vscode'
-import { pickFile, pickResources } from './resourcePicker'
+import { pickFile, pickResources, pickFiles } from './resourcePicker'
 import { Title } from './title'
 
 jest.mock('vscode')
@@ -67,3 +67,16 @@ describe('pickResources', () => {
     expect(pickedResources).toStrictEqual([mockedUri])
   })
 })
+
+describe('pickFiles', () => {
+  it('should return an array of selected file paths', async () => {
+    const files = [resolve('mock', 'file'), resolve('mock', 'file2')]
+    const mockedPickedUris = files.map(file => Uri.file(file))
+    const mockedTitle = 'insert great title here' as Title
+    mockedShowOpenDialog.mockResolvedValueOnce(mockedPickedUris)
+
+    const pickedResources = await pickFiles(mockedTitle)
+
+    expect(pickedResources).toStrictEqual(files)
+  })
+})
diff --git a/extension/src/vscode/resourcePicker.ts b/extension/src/vscode/resourcePicker.ts
index 7d47efb25f..651a6b98f7 100644
--- a/extension/src/vscode/resourcePicker.ts
+++ b/extension/src/vscode/resourcePicker.ts
@@ -5,19 +5,13 @@ export const pickFile = async (
   title: Title,
   filters?: OpenDialogOptions['filters']
 ): Promise<string | undefined> => {
-  const opts: OpenDialogOptions = {
+  const uris = await window.showOpenDialog({
     canSelectFolders: false,
     canSelectMany: false,
     filters,
     openLabel: 'Select',
     title
-  }
-
-  if (filters) {
-    opts.filters = filters
-  }
-
-  const uris = await window.showOpenDialog(opts)
+  })
 
   if (uris) {
     const [{ fsPath }] = uris

From a6b4a91b9ff252a3e3ecfc9b7d9c8074184cdce0 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 3 Oct 2023 08:25:36 -0500
Subject: [PATCH 06/35] Adjust how we handle failed files

---
 extension/src/fileSystem/index.test.ts      | 19 ++---------
 extension/src/fileSystem/index.ts           |  4 +--
 extension/src/pipeline/quickPick.test.ts    | 37 ++++++---------------
 extension/src/pipeline/quickPick.ts         |  2 +-
 extension/src/vscode/resourcePicker.test.ts | 10 +++---
 5 files changed, 22 insertions(+), 50 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 5e4bc8c2d5..90fa5be5ee 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -201,22 +201,7 @@ describe('loadDataFiles', () => {
       .mockReturnValueOnce(mockJsonContent)
 
     const resultWithErr = await loadDataFiles(dataFiles)
-    expect(resultWithErr).toStrictEqual([
-      {
-        data: [
-          { acc: 0.69, epoch: 10 },
-          { acc: 0.345, epoch: 11 }
-        ],
-        file: 'values.csv'
-      },
-      {
-        data: [
-          { acc: 0.69, epoch: 10 },
-          { acc: 0.345, epoch: 11 }
-        ],
-        file: 'file.json'
-      }
-    ])
+    expect(resultWithErr).toStrictEqual(undefined)
   })
 })
 
@@ -632,7 +617,7 @@ describe('addPlotToDvcYamlFile', () => {
     )
   })
 
-  it('should add the new plot with fields coming from different fiels', () => {
+  it('should add the new plot with fields coming from different files', () => {
     const mockDvcYamlContent = mockStagesLines.join('\n')
     const mockPlotYamlContent = [
       '',
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index f6bf8e50c3..4e32be1f95 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -355,12 +355,12 @@ export const loadDataFiles = async (
     const data = await loadDataFile(file)
 
     if (!data) {
-      continue
+      return undefined
     }
 
     filesData.push({ data, file })
   }
-  return filesData.length === 0 ? undefined : filesData
+  return filesData
 }
 
 export const writeJson = <
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 900a4c9cc7..d57cd5ef41 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -69,7 +69,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
   })
 
-  it('should show a toast message if files(s) fail to parse', async () => {
+  it('should show a toast message if the file or files fail to parse', async () => {
     mockedPickFiles.mockResolvedValueOnce(['data.csv'])
     mockedLoadDataFiles.mockResolvedValueOnce(undefined)
 
@@ -82,7 +82,7 @@ describe('pickPlotConfiguration', () => {
     )
   })
 
-  it('should show a toast message if an array of objects (with atleast two keys total) are not found within a file', async () => {
+  it('should show a toast message if an array of objects (with atleast two keys) are not found within a single chosen file', async () => {
     mockedPickFiles.mockResolvedValue(['file.yaml'])
     const invalidValues: unknown[] = [
       'string',
@@ -133,16 +133,16 @@ describe('pickPlotConfiguration', () => {
     }
   })
 
-  it('should show a toast message if atleast two fields are not found with multiple files', async () => {
+  it('should show a toast message if an array of objects (with atleast one key) are not found within multiple chosen files', async () => {
     mockedPickFiles.mockResolvedValueOnce([
       'file.yaml',
       'file2.yaml',
       'file3.yaml'
     ])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: '', file: 'file.yaml' },
-      { data: { val: [{ field1: 'only one field' }] }, file: 'file2.yaml' },
-      { data: [], file: 'file3.yaml' }
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' },
+      { data: [], file: 'file3.yaml' },
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' }
     ])
 
     const result = await pickPlotConfiguration()
@@ -183,36 +183,21 @@ describe('pickPlotConfiguration', () => {
     }
   })
 
-  it('should parse fields from a mixture of valid and invalid files if atleast 2 fields are found in total', async () => {
+  it('should parse fields from multiple valid data files (if atleast two fields are found total)', async () => {
     mockedPickFiles.mockResolvedValueOnce([
       'file.yaml',
       'file2.yaml',
       'file3.yaml'
     ])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      {
-        data: [
-          { field1: 1, field2: 2, field3: 1, field4: 2 },
-          { field1: 1, field2: 2, field3: 1, field4: 2 }
-        ],
-        file: 'file.yaml'
-      },
-      {
-        data: {
-          field1: [{ field1: 1 }, { field1: 3 }]
-        },
-        file: 'file3.yaml'
-      },
-      {
-        data: undefined,
-        file: 'file3.yaml'
-      }
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' },
+      { data: { val: [{ field1: 1 }] }, file: 'file2.yaml' }
     ])
 
-    await pickPlotConfiguration()
+    const result = await pickPlotConfiguration()
 
+    expect(result).toStrictEqual(undefined)
     expect(mockedShowError).not.toHaveBeenCalled()
-    expect(mockedQuickPickOne).toHaveBeenCalledTimes(1)
   })
 
   it('should let the user pick a template, x field, and y field', async () => {
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index c131003853..27172a1117 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -104,7 +104,7 @@ const getFieldOptionsFromArr = (
     const fields = getFieldOptions(data)
 
     if (fields.length === 0) {
-      continue
+      return { keys: {}, keysAmount: 0 }
     }
 
     keysAmount += fields.length
diff --git a/extension/src/vscode/resourcePicker.test.ts b/extension/src/vscode/resourcePicker.test.ts
index 38ae071e40..bd3f1b8625 100644
--- a/extension/src/vscode/resourcePicker.test.ts
+++ b/extension/src/vscode/resourcePicker.test.ts
@@ -70,13 +70,15 @@ describe('pickResources', () => {
 
 describe('pickFiles', () => {
   it('should return an array of selected file paths', async () => {
-    const files = [resolve('mock', 'file'), resolve('mock', 'file2')]
-    const mockedPickedUris = files.map(file => Uri.file(file))
+    const uris = [
+      Uri.file(resolve('mock', 'file')),
+      Uri.file(resolve('mock', 'file2'))
+    ]
     const mockedTitle = 'insert great title here' as Title
-    mockedShowOpenDialog.mockResolvedValueOnce(mockedPickedUris)
+    mockedShowOpenDialog.mockResolvedValueOnce(uris)
 
     const pickedResources = await pickFiles(mockedTitle)
 
-    expect(pickedResources).toStrictEqual(files)
+    expect(pickedResources).toStrictEqual(uris.map(uri => uri.fsPath))
   })
 })

From 3d405fb43d760615e2a4f9e2c44bee3473c91c7f Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 3 Oct 2023 09:17:23 -0500
Subject: [PATCH 07/35] Improve error handling

---
 extension/src/pipeline/quickPick.test.ts | 10 +++++++
 extension/src/pipeline/quickPick.ts      | 37 +++++++++++++++++-------
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index d57cd5ef41..c985c37c68 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -69,6 +69,16 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
   })
 
+  it('should show a toast message if the files are not the same data type', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['file.json', 'file.csv'])
+
+    const result = await pickPlotConfiguration()
+
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+    expect(mockedShowError).toHaveBeenCalledWith('Files must of the same type.')
+  })
+
   it('should show a toast message if the file or files fail to parse', async () => {
     mockedPickFiles.mockResolvedValueOnce(['data.csv'])
     mockedLoadDataFiles.mockResolvedValueOnce(undefined)
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 27172a1117..e64f8d921e 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -6,7 +6,7 @@ import {
   ValueTree,
   isValueTree
 } from '../cli/dvc/contract'
-import { loadDataFiles } from '../fileSystem'
+import { getFileExtension, loadDataFiles } from '../fileSystem'
 import { quickPickOne, quickPickValue } from '../vscode/quickPick'
 import { pickFiles } from '../vscode/resourcePicker'
 import { Title } from '../vscode/title'
@@ -78,7 +78,7 @@ const getFieldsFromArr = (dataArr: UnknownValue[]) => {
   return objsHaveSameKeys ? fieldObjKeys : []
 }
 
-const getFieldOptions = (data: UnknownValue): string[] => {
+const getFieldsFromValue = (data: UnknownValue): string[] => {
   const isArray = Array.isArray(data)
   const isObj = isValueTree(data)
   if (!isArray && !isObj) {
@@ -92,7 +92,12 @@ const getFieldOptions = (data: UnknownValue): string[] => {
     : []
 }
 
-const getFieldOptionsFromArr = (
+const showNotEnoughKeysToast = () =>
+  Toast.showError(
+    'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+  )
+
+const getFieldsFromDataFiles = (
   dataArr: { data: UnknownValue; file: string }[]
 ) => {
   const keys: {
@@ -101,10 +106,11 @@ const getFieldOptionsFromArr = (
   let keysAmount = 0
 
   for (const { file, data } of dataArr) {
-    const fields = getFieldOptions(data)
+    const fields = getFieldsFromValue(data)
 
     if (fields.length === 0) {
-      return { keys: {}, keysAmount: 0 }
+      void showNotEnoughKeysToast()
+      return
     }
 
     keysAmount += fields.length
@@ -115,7 +121,12 @@ const getFieldOptionsFromArr = (
     keys[file].push(...fields)
   }
 
-  return { keys, keysAmount }
+  if (keysAmount < 2) {
+    void showNotEnoughKeysToast()
+    return
+  }
+
+  return keys
 }
 
 export const pickPlotConfiguration = async (): Promise<
@@ -127,6 +138,12 @@ export const pickPlotConfiguration = async (): Promise<
     return
   }
 
+  const fileExts = new Set(files.map(file => getFileExtension(file)))
+
+  if (fileExts.size > 1) {
+    return Toast.showError('Files must of the same type.')
+  }
+
   const filesData = await loadDataFiles(files)
 
   if (!filesData) {
@@ -135,14 +152,12 @@ export const pickPlotConfiguration = async (): Promise<
     )
   }
 
-  const { keys, keysAmount } = getFieldOptionsFromArr(
+  const keys = getFieldsFromDataFiles(
     filesData as { data: UnknownValue; file: string }[]
   )
 
-  if (keysAmount < 2) {
-    return Toast.showError(
-      'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
-    )
+  if (!keys) {
+    return
   }
 
   const templateAndFields = await pickTemplateAndFields(keys)

From 76df14f39d972fa6cb05393ea0c71e1dada4ab17 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 3 Oct 2023 09:37:46 -0500
Subject: [PATCH 08/35] Fix broken pickFiles

---
 extension/src/vscode/resourcePicker.test.ts | 30 +++++++++++++++-----
 extension/src/vscode/resourcePicker.ts      | 31 ++++++++++++---------
 2 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/extension/src/vscode/resourcePicker.test.ts b/extension/src/vscode/resourcePicker.test.ts
index bd3f1b8625..b6b61ce2b1 100644
--- a/extension/src/vscode/resourcePicker.test.ts
+++ b/extension/src/vscode/resourcePicker.test.ts
@@ -69,16 +69,32 @@ describe('pickResources', () => {
 })
 
 describe('pickFiles', () => {
-  it('should return an array of selected file paths', async () => {
-    const uris = [
-      Uri.file(resolve('mock', 'file')),
-      Uri.file(resolve('mock', 'file2'))
+  it('should call window.showOpenDialog with the correct options', async () => {
+    const mockedTitle = 'I decided to not decide' as Title
+    mockedShowOpenDialog.mockResolvedValueOnce(undefined)
+
+    await pickFiles(mockedTitle, { Images: ['png', 'jpg'] })
+
+    expect(mockedShowOpenDialog).toHaveBeenCalledWith({
+      canSelectFiles: true,
+      canSelectFolders: false,
+      canSelectMany: true,
+      filters: { Images: ['png', 'jpg'] },
+      openLabel: 'Select',
+      title: mockedTitle
+    })
+  })
+
+  it('should return an array of paths if any are selected', async () => {
+    const mockedUris = [
+      Uri.file(resolve('mock', 'file1.json')),
+      Uri.file(resolve('mock', 'file2.json'))
     ]
-    const mockedTitle = 'insert great title here' as Title
-    mockedShowOpenDialog.mockResolvedValueOnce(uris)
+    const mockedTitle = 'this is even more fun' as Title
+    mockedShowOpenDialog.mockResolvedValueOnce(mockedUris)
 
     const pickedResources = await pickFiles(mockedTitle)
 
-    expect(pickedResources).toStrictEqual(uris.map(uri => uri.fsPath))
+    expect(pickedResources).toStrictEqual(mockedUris.map(uri => uri.fsPath))
   })
 })
diff --git a/extension/src/vscode/resourcePicker.ts b/extension/src/vscode/resourcePicker.ts
index 651a6b98f7..d5fc137505 100644
--- a/extension/src/vscode/resourcePicker.ts
+++ b/extension/src/vscode/resourcePicker.ts
@@ -19,30 +19,35 @@ export const pickFile = async (
   }
 }
 
-export const pickResources = (
+export const pickResources = (title: Title): Thenable<Uri[] | undefined> => {
+  return window.showOpenDialog({
+    canSelectFiles: true,
+    canSelectFolders: true,
+    canSelectMany: true,
+    openLabel: 'Select',
+    title
+  })
+}
+
+export const pickFiles = async (
   title: Title,
   filters?: OpenDialogOptions['filters']
-): Thenable<Uri[] | undefined> => {
+): Promise<string[] | undefined> => {
   const opts: OpenDialogOptions = {
     canSelectFiles: true,
-    canSelectFolders: true,
+    canSelectFolders: false,
     canSelectMany: true,
     openLabel: 'Select',
     title
   }
 
-  if (!filters) {
+  if (filters) {
     opts.filters = filters
   }
 
-  return window.showOpenDialog(opts)
-}
-
-export const pickFiles = async (
-  title: Title,
-  filters?: OpenDialogOptions['filters']
-): Promise<string[] | undefined> => {
-  const files = await pickResources(title, filters)
+  const uris = await window.showOpenDialog(opts)
 
-  return files?.map(({ fsPath }) => fsPath)
+  if (uris) {
+    return uris.map(({ fsPath }) => fsPath)
+  }
 }

From 22c64c50477df429e888f33e48b21e861d2e26a2 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 3 Oct 2023 10:39:14 -0500
Subject: [PATCH 09/35] Use relative path instead of just filename in quick
 pick

---
 extension/src/pipeline/index.ts          |  2 +-
 extension/src/pipeline/quickPick.test.ts | 72 ++++++++++++------------
 extension/src/pipeline/quickPick.ts      | 21 ++++---
 3 files changed, 49 insertions(+), 46 deletions(-)

diff --git a/extension/src/pipeline/index.ts b/extension/src/pipeline/index.ts
index d69f0944f8..ae019ad728 100644
--- a/extension/src/pipeline/index.ts
+++ b/extension/src/pipeline/index.ts
@@ -119,7 +119,7 @@ export class Pipeline extends DeferredDisposable {
       return
     }
 
-    const plotConfiguration = await pickPlotConfiguration()
+    const plotConfiguration = await pickPlotConfiguration(cwd)
 
     if (!plotConfiguration) {
       return
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index c985c37c68..0c08efe5ee 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -54,7 +54,7 @@ describe('pickPlotConfiguration', () => {
       { data: mockValidData, file: 'file.json' }
     ])
 
-    await pickPlotConfiguration()
+    await pickPlotConfiguration('/')
 
     expect(mockedPickFiles).toHaveBeenCalledWith(Title.SELECT_PLOT_DATA, {
       'Data Formats': ['json', 'csv', 'tsv', 'yaml']
@@ -64,7 +64,7 @@ describe('pickPlotConfiguration', () => {
   it('should return early if user does not select a file', async () => {
     mockedPickFiles.mockResolvedValueOnce(undefined)
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(result).toStrictEqual(undefined)
   })
@@ -72,7 +72,7 @@ describe('pickPlotConfiguration', () => {
   it('should show a toast message if the files are not the same data type', async () => {
     mockedPickFiles.mockResolvedValueOnce(['file.json', 'file.csv'])
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
@@ -83,7 +83,7 @@ describe('pickPlotConfiguration', () => {
     mockedPickFiles.mockResolvedValueOnce(['data.csv'])
     mockedLoadDataFiles.mockResolvedValueOnce(undefined)
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
@@ -133,7 +133,7 @@ describe('pickPlotConfiguration', () => {
         { data: invalidVal, file: 'file.yaml' }
       ])
 
-      result = await pickPlotConfiguration()
+      result = await pickPlotConfiguration('/')
 
       expect(result).toStrictEqual(undefined)
       expect(mockedShowError).toHaveBeenCalledTimes(1 + ind)
@@ -155,7 +155,7 @@ describe('pickPlotConfiguration', () => {
       { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' }
     ])
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
@@ -186,7 +186,7 @@ describe('pickPlotConfiguration', () => {
         { data: val, file: 'file.yaml' }
       ])
 
-      await pickPlotConfiguration()
+      await pickPlotConfiguration('/')
 
       expect(mockedShowError).not.toHaveBeenCalled()
       expect(mockedQuickPickOne).toHaveBeenCalledTimes(ind + 1)
@@ -204,23 +204,23 @@ describe('pickPlotConfiguration', () => {
       { data: { val: [{ field1: 1 }] }, file: 'file2.yaml' }
     ])
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).not.toHaveBeenCalled()
   })
 
   it('should let the user pick a template, x field, and y field', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: mockValidData, file: 'file.json' }
+      { data: mockValidData, file: '/file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedQuickPickValue
-      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
-      .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
+      .mockResolvedValueOnce({ file: '/file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: '/file.json', key: 'prob' })
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(mockedQuickPickOne).toHaveBeenNthCalledWith(
       1,
@@ -245,8 +245,8 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
+        { label: 'actual', value: { file: '/file.json', key: 'actual' } },
+        { label: 'prob', value: { file: '/file.json', key: 'prob' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -258,7 +258,7 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
+        { label: 'prob', value: { file: '/file.json', key: 'prob' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -266,23 +266,23 @@ describe('pickPlotConfiguration', () => {
     )
     expect(result).toStrictEqual({
       template: 'simple',
-      x: { file: 'file.json', key: 'actual' },
-      y: { file: 'file.json', key: 'prob' }
+      x: { file: '/file.json', key: 'actual' },
+      y: { file: '/file.json', key: 'prob' }
     })
   })
 
   it('should let the user pick a x field and y field from multiple files', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['file.json', 'file2.json'])
+    mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file2.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: mockValidData, file: 'file.json' },
-      { data: mockValidData, file: 'file2.json' }
+      { data: mockValidData, file: '/file.json' },
+      { data: mockValidData, file: '/file2.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedQuickPickValue
-      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
-      .mockResolvedValueOnce({ file: 'file2.json', key: 'prob' })
+      .mockResolvedValueOnce({ file: '/file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: '/file2.json', key: 'prob' })
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
       1,
@@ -292,15 +292,15 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        { label: 'actual', value: { file: '/file.json', key: 'actual' } },
+        { label: 'prob', value: { file: '/file.json', key: 'prob' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'actual', value: { file: '/file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: '/file2.json', key: 'prob' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -312,14 +312,14 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        { label: 'prob', value: { file: '/file.json', key: 'prob' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'actual', value: { file: '/file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: '/file2.json', key: 'prob' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -327,8 +327,8 @@ describe('pickPlotConfiguration', () => {
     )
     expect(result).toStrictEqual({
       template: 'simple',
-      x: { file: 'file.json', key: 'actual' },
-      y: { file: 'file2.json', key: 'prob' }
+      x: { file: '/file.json', key: 'actual' },
+      y: { file: '/file2.json', key: 'prob' }
     })
   })
 
@@ -339,7 +339,7 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce(undefined)
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(mockedQuickPickOne).toHaveBeenCalledTimes(1)
 
@@ -354,7 +354,7 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedQuickPickValue.mockResolvedValueOnce(undefined)
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(mockedQuickPickValue).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
@@ -370,7 +370,7 @@ describe('pickPlotConfiguration', () => {
       .mockResolvedValueOnce('actual')
       .mockResolvedValueOnce(undefined)
 
-    const result = await pickPlotConfiguration()
+    const result = await pickPlotConfiguration('/')
 
     expect(mockedQuickPickValue).toHaveBeenCalledTimes(2)
     expect(result).toStrictEqual(undefined)
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index e64f8d921e..34a89d2fc7 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -1,3 +1,4 @@
+import { relative } from 'path'
 import isEqual from 'lodash.isequal'
 import { QuickPickItemKind } from 'vscode'
 import {
@@ -11,16 +12,18 @@ import { quickPickOne, quickPickValue } from '../vscode/quickPick'
 import { pickFiles } from '../vscode/resourcePicker'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
-import { getFileName } from '../fileSystem/util'
 
 const pickDataFiles = (): Thenable<string[] | undefined> =>
   pickFiles(Title.SELECT_PLOT_DATA, {
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
 
-const pickTemplateAndFields = async (fields: {
-  [file: string]: string[]
-}): Promise<PlotConfigData | undefined> => {
+const pickTemplateAndFields = async (
+  cwd: string,
+  fields: {
+    [file: string]: string[]
+  }
+): Promise<PlotConfigData | undefined> => {
   const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
 
   if (!template) {
@@ -33,7 +36,7 @@ const pickTemplateAndFields = async (fields: {
     items.push(
       {
         kind: QuickPickItemKind.Separator,
-        label: getFileName(file),
+        label: relative(cwd, file),
         value: undefined
       },
       ...keys.map(key => ({ label: key, value: { file, key } }))
@@ -129,9 +132,9 @@ const getFieldsFromDataFiles = (
   return keys
 }
 
-export const pickPlotConfiguration = async (): Promise<
-  PlotConfigData | undefined
-> => {
+export const pickPlotConfiguration = async (
+  cwd: string
+): Promise<PlotConfigData | undefined> => {
   const files = await pickDataFiles()
 
   if (!files) {
@@ -160,7 +163,7 @@ export const pickPlotConfiguration = async (): Promise<
     return
   }
 
-  const templateAndFields = await pickTemplateAndFields(keys)
+  const templateAndFields = await pickTemplateAndFields(cwd, keys)
 
   if (!templateAndFields) {
     return

From da06e65a0273232caeafa99e3bdf14be1f6240f0 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 3 Oct 2023 11:05:26 -0500
Subject: [PATCH 10/35] Delete unneeded function

---
 extension/src/fileSystem/util.ts | 2 --
 1 file changed, 2 deletions(-)

diff --git a/extension/src/fileSystem/util.ts b/extension/src/fileSystem/util.ts
index 53effd7ac0..36ab159925 100644
--- a/extension/src/fileSystem/util.ts
+++ b/extension/src/fileSystem/util.ts
@@ -22,6 +22,4 @@ export const getParent = (pathArray: string[], idx: number) => {
 export const removeTrailingSlash = (path: string): string =>
   path.endsWith(sep) ? path.slice(0, -1) : path
 
-export const getFileName = (path: string) => parse(path).base
-
 export const getFileNameWithoutExt = (path: string) => parse(path).name

From b6af87c6aa4a837b9be0d3ecf9e3a706fc35d279 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 4 Oct 2023 08:36:07 -0500
Subject: [PATCH 11/35] Resolve some comments: * simplify plot object *
 simplify resourcePicker

---
 extension/src/fileSystem/index.test.ts      | 10 +++--
 extension/src/fileSystem/index.ts           |  9 +---
 extension/src/vscode/resourcePicker.test.ts |  1 +
 extension/src/vscode/resourcePicker.ts      | 49 ++++++++-------------
 4 files changed, 27 insertions(+), 42 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 90fa5be5ee..65a5874bc4 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -592,10 +592,11 @@ describe('addPlotToDvcYamlFile', () => {
     '        eval/prc/test.json: precision'
   ]
   const mockNewPlotLines = [
-    '  - data.json:',
+    '  - simple_plot:',
     '      template: simple',
     '      x: epochs',
-    '      y: accuracy'
+    '      y:',
+    '        data.json: accuracy'
   ]
   it('should add a plots list with the new plot if the dvc.yaml file has no plots', () => {
     const mockDvcYamlContent = mockStagesLines.join('\n')
@@ -700,10 +701,11 @@ describe('addPlotToDvcYamlFile', () => {
     ].join('\n')
     const mockPlotYamlContent = [
       '',
-      '    - data.json:',
+      '    - simple_plot:',
       '          template: simple',
       '          x: epochs',
-      '          y: accuracy',
+      '          y:',
+      '              data.json: accuracy',
       ''
     ].join('\n')
     mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index 4e32be1f95..f79dcb6211 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -216,18 +216,11 @@ const loadYamlAsDoc = (
 
 const getPlotYamlObj = (cwd: string, plot: PlotConfigData) => {
   const { x, y, template } = plot
-  const usesSingleFile = x.file === y.file
-  if (usesSingleFile) {
-    const dataFile = x.file
-    const plotName = relative(cwd, dataFile)
-    return { [plotName]: { template, x: x.key, y: y.key } }
-  }
-
   const plotName = `${template}_plot`
   return {
     [plotName]: {
       template,
-      x: { [relative(cwd, x.file)]: x.key },
+      x: x.file === y.file ? x.key : { [relative(cwd, x.file)]: x.key },
       y: { [relative(cwd, y.file)]: y.key }
     }
   }
diff --git a/extension/src/vscode/resourcePicker.test.ts b/extension/src/vscode/resourcePicker.test.ts
index b6b61ce2b1..69050eb2dc 100644
--- a/extension/src/vscode/resourcePicker.test.ts
+++ b/extension/src/vscode/resourcePicker.test.ts
@@ -21,6 +21,7 @@ describe('pickFile', () => {
     await pickFile(mockedTitle)
 
     expect(mockedShowOpenDialog).toHaveBeenCalledWith({
+      canSelectFiles: true,
       canSelectFolders: false,
       canSelectMany: false,
       openLabel: 'Select',
diff --git a/extension/src/vscode/resourcePicker.ts b/extension/src/vscode/resourcePicker.ts
index d5fc137505..87991ee764 100644
--- a/extension/src/vscode/resourcePicker.ts
+++ b/extension/src/vscode/resourcePicker.ts
@@ -1,51 +1,40 @@
 import { OpenDialogOptions, Uri, window } from 'vscode'
 import { Title } from './title'
 
-export const pickFile = async (
+export const pickResources = (
   title: Title,
+  canSelectFolders = true,
+  canSelectMany = true,
   filters?: OpenDialogOptions['filters']
-): Promise<string | undefined> => {
-  const uris = await window.showOpenDialog({
-    canSelectFolders: false,
-    canSelectMany: false,
-    filters,
+): Thenable<Uri[] | undefined> => {
+  const opts: OpenDialogOptions = {
+    canSelectFiles: true,
+    canSelectFolders,
+    canSelectMany,
     openLabel: 'Select',
     title
-  })
+  }
+
+  if (filters) {
+    opts.filters = filters
+  }
+
+  return window.showOpenDialog(opts)
+}
 
+export const pickFile = async (title: Title): Promise<string | undefined> => {
+  const uris = await pickResources(title, false, false)
   if (uris) {
     const [{ fsPath }] = uris
     return fsPath
   }
 }
 
-export const pickResources = (title: Title): Thenable<Uri[] | undefined> => {
-  return window.showOpenDialog({
-    canSelectFiles: true,
-    canSelectFolders: true,
-    canSelectMany: true,
-    openLabel: 'Select',
-    title
-  })
-}
-
 export const pickFiles = async (
   title: Title,
   filters?: OpenDialogOptions['filters']
 ): Promise<string[] | undefined> => {
-  const opts: OpenDialogOptions = {
-    canSelectFiles: true,
-    canSelectFolders: false,
-    canSelectMany: true,
-    openLabel: 'Select',
-    title
-  }
-
-  if (filters) {
-    opts.filters = filters
-  }
-
-  const uris = await window.showOpenDialog(opts)
+  const uris = await pickResources(title, false, true, filters)
 
   if (uris) {
     return uris.map(({ fsPath }) => fsPath)

From bd3e211c441480c7c620df19d4b03d13ed039c32 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 4 Oct 2023 09:25:02 -0500
Subject: [PATCH 12/35] Improve text

---
 extension/src/pipeline/quickPick.test.ts | 6 +++---
 extension/src/pipeline/quickPick.ts      | 4 ++--
 extension/src/plots/quickPick.test.ts    | 2 +-
 extension/src/plots/quickPick.ts         | 2 +-
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 0c08efe5ee..05bf97bfd2 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -88,7 +88,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'Failed to parse the requested file(s). Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'Failed to parse the requested file or files. Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
@@ -138,7 +138,7 @@ describe('pickPlotConfiguration', () => {
       expect(result).toStrictEqual(undefined)
       expect(mockedShowError).toHaveBeenCalledTimes(1 + ind)
       expect(mockedShowError).toHaveBeenCalledWith(
-        'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+        'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
       )
     }
   })
@@ -160,7 +160,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 34a89d2fc7..0a2d994431 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -97,7 +97,7 @@ const getFieldsFromValue = (data: UnknownValue): string[] => {
 
 const showNotEnoughKeysToast = () =>
   Toast.showError(
-    'The requested file(s) does not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+    'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
   )
 
 const getFieldsFromDataFiles = (
@@ -151,7 +151,7 @@ export const pickPlotConfiguration = async (
 
   if (!filesData) {
     return Toast.showError(
-      'Failed to parse the requested file(s). Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'Failed to parse the requested file or files. Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   }
 
diff --git a/extension/src/plots/quickPick.test.ts b/extension/src/plots/quickPick.test.ts
index 7e22b95dee..8bcaa48235 100644
--- a/extension/src/plots/quickPick.test.ts
+++ b/extension/src/plots/quickPick.test.ts
@@ -20,7 +20,7 @@ describe('pickPlotType', () => {
       [
         {
           description:
-            'Create a data series plot definition by selecting data from a file',
+            'Create a data series plot definition by selecting data from one or more files (To select multiple files, hold the Ctrl/Command key and click on the files.).',
           label: 'Data Series',
           value: 'data-series'
         },
diff --git a/extension/src/plots/quickPick.ts b/extension/src/plots/quickPick.ts
index 5c00508247..e2358bb262 100644
--- a/extension/src/plots/quickPick.ts
+++ b/extension/src/plots/quickPick.ts
@@ -11,7 +11,7 @@ export const pickPlotType = () =>
     [
       {
         description:
-          'Create a data series plot definition by selecting data from a file',
+          'Create a data series plot definition by selecting data from one or more files (To select multiple files, hold the Ctrl/Command key and click on the files.).',
         label: 'Data Series',
         value: PLOT_TYPE.DATA_SERIES
       },

From 52defff3defdf4401936e92382d96e16b2919452 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 4 Oct 2023 18:09:07 -0500
Subject: [PATCH 13/35] wip solution

---
 extension/src/fileSystem/index.ts   |   9 +-
 extension/src/pipeline/quickPick.ts | 180 ++++++++++++++++++++--------
 2 files changed, 133 insertions(+), 56 deletions(-)

diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index f79dcb6211..a7d41d4b69 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -342,17 +342,14 @@ const loadDataFile = (file: string): unknown => {
 
 export const loadDataFiles = async (
   files: string[]
-): Promise<{ file: string; data: unknown }[] | undefined> => {
+): Promise<{ file: string; data: unknown }[]> => {
   const filesData: { file: string; data: unknown }[] = []
+
   for (const file of files) {
     const data = await loadDataFile(file)
-
-    if (!data) {
-      return undefined
-    }
-
     filesData.push({ data, file })
   }
+
   return filesData
 }
 
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 0a2d994431..190df91f53 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -13,6 +13,16 @@ import { pickFiles } from '../vscode/resourcePicker'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
 
+export type PlotConfigData = {
+  x: { file: string; key: string }
+  template: string
+  y: { file: string; key: string }
+}
+
+type UnknownValue = Value | ValueTree
+
+type fileFields = { [file: string]: string[] }
+
 const pickDataFiles = (): Thenable<string[] | undefined> =>
   pickFiles(Title.SELECT_PLOT_DATA, {
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
@@ -20,9 +30,7 @@ const pickDataFiles = (): Thenable<string[] | undefined> =>
 
 const pickTemplateAndFields = async (
   cwd: string,
-  fields: {
-    [file: string]: string[]
-  }
+  fields: fileFields
 ): Promise<PlotConfigData | undefined> => {
   const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
 
@@ -61,109 +69,181 @@ const pickTemplateAndFields = async (
   return { template, x, y }
 }
 
-export type PlotConfigData = {
-  x: { file: string; key: string }
-  template: string
-  y: { file: string; key: string }
+const joinList = (items: string[]) => {
+  if (items.length <= 2) {
+    return items.join(' and ')
+  }
+
+  return `${[items.slice(0, -1)].join(', ')}, and ${items[items.length - 1]}`
 }
 
-type UnknownValue = Value | ValueTree
+const validateFileNames = (files: string[] | undefined) => {
+  if (!files) {
+    return []
+  }
+
+  const fileExts = [...new Set(files.map(file => getFileExtension(file)))]
+
+  if (fileExts.length > 1) {
+    void Toast.showError(
+      `Found files with ${joinList(
+        fileExts
+      )} extensions. Files must be of the same type.`
+    )
+    return files
+  }
+  return files
+}
 
-const getFieldsFromArr = (dataArr: UnknownValue[]) => {
+const getFieldsFromArr = (
+  dataArr: UnknownValue[]
+): { arrLength: number; fields: string[] } | undefined => {
   const firstArrVal: UnknownValue = dataArr[0]
   if (!isValueTree(firstArrVal)) {
-    return []
+    return
   }
   const fieldObjKeys = Object.keys(firstArrVal)
   const objsHaveSameKeys = dataArr.every(
     val => isValueTree(val) && isEqual(fieldObjKeys, Object.keys(val))
   )
-  return objsHaveSameKeys ? fieldObjKeys : []
+  if (!objsHaveSameKeys) {
+    return
+  }
+  return { arrLength: dataArr.length, fields: fieldObjKeys }
 }
 
-const getFieldsFromValue = (data: UnknownValue): string[] => {
+const getFieldsFromValue = (
+  data: UnknownValue
+): { arrLength: number; fields: string[] } | undefined => {
   const isArray = Array.isArray(data)
   const isObj = isValueTree(data)
   if (!isArray && !isObj) {
-    return []
+    return
   }
 
   const maybeFieldsObjArr = isArray ? data : data[Object.keys(data)[0]]
 
-  return Array.isArray(maybeFieldsObjArr)
-    ? getFieldsFromArr(maybeFieldsObjArr)
-    : []
+  if (!Array.isArray(maybeFieldsObjArr)) {
+    return
+  }
+
+  return getFieldsFromArr(maybeFieldsObjArr)
 }
 
-const showNotEnoughKeysToast = () =>
-  Toast.showError(
-    'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+const showNotEnoughKeysToast = (files: string[]) => {
+  const isSingle = files.length === 1
+  return Toast.showError(
+    `${joinList(files)} ${
+      isSingle ? 'does' : 'do'
+    } not contain enough keys (columns) to generate a plot. Does the ${
+      isSingle ? 'file' : 'files'
+    } follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?`
   )
+}
+
+const validateSingleDataFileFields = ({
+  file,
+  data
+}: {
+  file: string
+  data: UnknownValue
+}) => {
+  const { fields = [] } = getFieldsFromValue(data) || {}
+
+  if (fields.length < 2) {
+    void showNotEnoughKeysToast([file])
+    return
+  }
+
+  return { [file]: fields }
+}
 
 const getFieldsFromDataFiles = (
   dataArr: { data: UnknownValue; file: string }[]
 ) => {
-  const keys: {
-    [file: string]: string[]
-  } = {}
-  let keysAmount = 0
+  const failedFiles: string[] = []
+  const filesArrLength: Set<number> = new Set()
+  const keys: fileFields = {}
 
   for (const { file, data } of dataArr) {
-    const fields = getFieldsFromValue(data)
+    const fileFields = getFieldsFromValue(data)
 
-    if (fields.length === 0) {
-      void showNotEnoughKeysToast()
-      return
+    if (!fileFields) {
+      failedFiles.push(file)
+      continue
     }
 
-    keysAmount += fields.length
+    const { fields, arrLength } = fileFields
 
     if (!keys[file]) {
       keys[file] = []
     }
+    filesArrLength.add(arrLength)
     keys[file].push(...fields)
   }
 
-  if (keysAmount < 2) {
-    void showNotEnoughKeysToast()
+  return { failedFiles, filesArrLength, keys }
+}
+
+const validateMultiDataFileFields = (
+  dataArr: { data: UnknownValue; file: string }[]
+) => {
+  const { keys, failedFiles, filesArrLength } = getFieldsFromDataFiles(dataArr)
+
+  if (failedFiles) {
+    void showNotEnoughKeysToast(failedFiles)
     return
   }
 
+  if (filesArrLength.size > 1) {
+    void Toast.showError(
+      'The files must have the same array (or row in tsv/csv) length.'
+    )
+  }
+
   return keys
 }
 
+const validateFilesData = async (files: string[]) => {
+  const filesData = (await loadDataFiles(files)) as {
+    data: UnknownValue
+    file: string
+  }[]
+
+  const failedFiles = filesData.filter(({ data }) => !!data)
+
+  if (failedFiles.length > 0) {
+    const files = joinList(failedFiles.map(({ file }) => file))
+    void Toast.showError(
+      `Failed to parse ${files}. ${
+        files.length === 1 ? 'Does the file' : 'Do the files'
+      } contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?`
+    )
+    return false
+  }
+
+  return files.length === 1
+    ? validateSingleDataFileFields(filesData[0])
+    : validateMultiDataFileFields(filesData)
+}
+
 export const pickPlotConfiguration = async (
   cwd: string
 ): Promise<PlotConfigData | undefined> => {
   const files = await pickDataFiles()
+  const validFileNames = validateFileNames(files)
 
-  if (!files) {
+  if (validFileNames.length === 0) {
     return
   }
 
-  const fileExts = new Set(files.map(file => getFileExtension(file)))
-
-  if (fileExts.size > 1) {
-    return Toast.showError('Files must of the same type.')
-  }
-
-  const filesData = await loadDataFiles(files)
-
-  if (!filesData) {
-    return Toast.showError(
-      'Failed to parse the requested file or files. Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
-    )
-  }
-
-  const keys = getFieldsFromDataFiles(
-    filesData as { data: UnknownValue; file: string }[]
-  )
+  const validFileFields = await validateFilesData(validFileNames)
 
-  if (!keys) {
+  if (!validFileFields) {
     return
   }
 
-  const templateAndFields = await pickTemplateAndFields(cwd, keys)
+  const templateAndFields = await pickTemplateAndFields(cwd, validFileFields)
 
   if (!templateAndFields) {
     return

From ec182c35076fdcd7ed6be69d2dd82a6e749a67e1 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 5 Oct 2023 09:03:32 -0500
Subject: [PATCH 14/35] fix broken tests

---
 extension/src/fileSystem/index.test.ts   | 40 +++++++----
 extension/src/fileSystem/index.ts        | 19 ++---
 extension/src/pipeline/quickPick.test.ts | 92 ++++++++++++------------
 extension/src/pipeline/quickPick.ts      | 41 +++++------
 4 files changed, 101 insertions(+), 91 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 65a5874bc4..4fc43a5145 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -182,7 +182,7 @@ describe('loadDataFiles', () => {
     for (const file of dataFiles) {
       const resultWithErr = await loadDataFiles([file])
 
-      expect(resultWithErr).toStrictEqual(undefined)
+      expect(resultWithErr).toStrictEqual([{ data: undefined, file }])
     }
   })
 
@@ -201,7 +201,23 @@ describe('loadDataFiles', () => {
       .mockReturnValueOnce(mockJsonContent)
 
     const resultWithErr = await loadDataFiles(dataFiles)
-    expect(resultWithErr).toStrictEqual(undefined)
+    expect(resultWithErr).toStrictEqual([
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'values.csv'
+      },
+      { data: undefined, file: 'file.tsv' },
+      {
+        data: [
+          { acc: 0.69, epoch: 10 },
+          { acc: 0.345, epoch: 11 }
+        ],
+        file: 'file.json'
+      }
+    ])
   })
 })
 
@@ -608,8 +624,8 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
-      x: { file: '/data.json', key: 'epochs' },
-      y: { file: '/data.json', key: 'accuracy' }
+      x: { file: 'data.json', key: 'epochs' },
+      y: { file: 'data.json', key: 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -636,8 +652,8 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
-      x: { file: '/data.json', key: 'epochs' },
-      y: { file: '/acc.json', key: 'accuracy' }
+      x: { file: 'data.json', key: 'epochs' },
+      y: { file: 'acc.json', key: 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -654,8 +670,8 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
-      x: { file: '/data.json', key: 'epochs' },
-      y: { file: '/data.json', key: 'accuracy' }
+      x: { file: 'data.json', key: 'epochs' },
+      y: { file: 'data.json', key: 'accuracy' }
     })
 
     mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
@@ -676,8 +692,8 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
-      x: { file: '/data.json', key: 'epochs' },
-      y: { file: '/data.json', key: 'accuracy' }
+      x: { file: 'data.json', key: 'epochs' },
+      y: { file: 'data.json', key: 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -713,8 +729,8 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
-      x: { file: '/data.json', key: 'epochs' },
-      y: { file: '/data.json', key: 'accuracy' }
+      x: { file: 'data.json', key: 'epochs' },
+      y: { file: 'data.json', key: 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index a7d41d4b69..c2ce48ce1b 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -214,23 +214,19 @@ const loadYamlAsDoc = (
   }
 }
 
-const getPlotYamlObj = (cwd: string, plot: PlotConfigData) => {
+const getPlotYamlObj = (plot: PlotConfigData) => {
   const { x, y, template } = plot
   const plotName = `${template}_plot`
   return {
     [plotName]: {
       template,
-      x: x.file === y.file ? x.key : { [relative(cwd, x.file)]: x.key },
-      y: { [relative(cwd, y.file)]: y.key }
+      x: x.file === y.file ? x.key : { [x.file]: x.key },
+      y: { [y.file]: y.key }
     }
   }
 }
 
-const getPlotsYaml = (
-  cwd: string,
-  plotObj: PlotConfigData,
-  indentSearchLines: string[]
-) => {
+const getPlotsYaml = (plotObj: PlotConfigData, indentSearchLines: string[]) => {
   const indentReg = /^( +)[^ ]/
   const indentLine = indentSearchLines.find(line => indentReg.test(line)) || ''
   const spacesMatches = indentLine.match(indentReg)
@@ -238,7 +234,7 @@ const getPlotsYaml = (
 
   return yaml
     .stringify(
-      { plots: [getPlotYamlObj(cwd, plotObj)] },
+      { plots: [getPlotYamlObj(plotObj)] },
       { indent: spaces ? spaces.length : 2 }
     )
     .split('\n')
@@ -258,7 +254,7 @@ export const addPlotToDvcYamlFile = (cwd: string, plotObj: PlotConfigData) => {
   const plots = doc.get('plots', true) as yaml.YAMLSeq | undefined
 
   if (!plots?.range) {
-    const plotYaml = getPlotsYaml(cwd, plotObj, dvcYamlLines)
+    const plotYaml = getPlotsYaml(plotObj, dvcYamlLines)
     dvcYamlLines.push(...plotYaml)
     writeFileSync(dvcYamlFile, dvcYamlLines.join('\n'))
     return
@@ -272,7 +268,6 @@ export const addPlotToDvcYamlFile = (cwd: string, plotObj: PlotConfigData) => {
 
   const plotsStartPos = lineCounter.linePos(plots.range[0]).line - 1
   const plotYaml = getPlotsYaml(
-    cwd,
     plotObj,
     dvcYamlLines.slice(plotsStartPos, insertLineNum)
   )
@@ -344,12 +339,10 @@ export const loadDataFiles = async (
   files: string[]
 ): Promise<{ file: string; data: unknown }[]> => {
   const filesData: { file: string; data: unknown }[] = []
-
   for (const file of files) {
     const data = await loadDataFile(file)
     filesData.push({ data, file })
   }
-
   return filesData
 }
 
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 05bf97bfd2..66d292761e 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -49,7 +49,7 @@ const mockValidData = [
 
 describe('pickPlotConfiguration', () => {
   it('should let the user pick from files with accepted data types', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
@@ -70,30 +70,34 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should show a toast message if the files are not the same data type', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['file.json', 'file.csv'])
+    mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file.csv'])
 
     const result = await pickPlotConfiguration('/')
 
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
-    expect(mockedShowError).toHaveBeenCalledWith('Files must of the same type.')
+    expect(mockedShowError).toHaveBeenCalledWith(
+      'Found files with .json and .csv extensions. Files must be of the same type.'
+    )
   })
 
   it('should show a toast message if the file or files fail to parse', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['data.csv'])
-    mockedLoadDataFiles.mockResolvedValueOnce(undefined)
+    mockedPickFiles.mockResolvedValueOnce(['/data.csv'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: undefined, file: '/data.csv' }
+    ])
 
     const result = await pickPlotConfiguration('/')
 
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'Failed to parse the requested file or files. Does the file or files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'Failed to parse data.csv. Do the files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
   it('should show a toast message if an array of objects (with atleast two keys) are not found within a single chosen file', async () => {
-    mockedPickFiles.mockResolvedValue(['file.yaml'])
+    mockedPickFiles.mockResolvedValue(['/file.yaml'])
     const invalidValues: unknown[] = [
       'string',
       13,
@@ -130,7 +134,7 @@ describe('pickPlotConfiguration', () => {
 
     for (const [ind, invalidVal] of invalidValues.entries()) {
       mockedLoadDataFiles.mockResolvedValueOnce([
-        { data: invalidVal, file: 'file.yaml' }
+        { data: invalidVal, file: '/file.yaml' }
       ])
 
       result = await pickPlotConfiguration('/')
@@ -138,21 +142,21 @@ describe('pickPlotConfiguration', () => {
       expect(result).toStrictEqual(undefined)
       expect(mockedShowError).toHaveBeenCalledTimes(1 + ind)
       expect(mockedShowError).toHaveBeenCalledWith(
-        'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+        'file.yaml does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
       )
     }
   })
 
   it('should show a toast message if an array of objects (with atleast one key) are not found within multiple chosen files', async () => {
     mockedPickFiles.mockResolvedValueOnce([
-      'file.yaml',
-      'file2.yaml',
-      'file3.yaml'
+      '/file.yaml',
+      '/file2.yaml',
+      '/file3.yaml'
     ])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' },
-      { data: [], file: 'file3.yaml' },
-      { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' }
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file2.yaml' },
+      { data: [], file: '/file3.yaml' },
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file2.yaml' }
     ])
 
     const result = await pickPlotConfiguration('/')
@@ -160,12 +164,12 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'The requested file or files do not contain enough keys (columns) to generate a plot. Does the file or files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'file3.yaml does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
   it('should parse fields from a valid data file', async () => {
-    mockedPickFiles.mockResolvedValue(['file.yaml'])
+    mockedPickFiles.mockResolvedValue(['/file.yaml'])
     const validValues: unknown[] = [
       [{ field1: 1, field2: 2 }],
       [
@@ -183,7 +187,7 @@ describe('pickPlotConfiguration', () => {
 
     for (const [ind, val] of validValues.entries()) {
       mockedLoadDataFiles.mockResolvedValueOnce([
-        { data: val, file: 'file.yaml' }
+        { data: val, file: '/file.yaml' }
       ])
 
       await pickPlotConfiguration('/')
@@ -194,14 +198,10 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should parse fields from multiple valid data files (if atleast two fields are found total)', async () => {
-    mockedPickFiles.mockResolvedValueOnce([
-      'file.yaml',
-      'file2.yaml',
-      'file3.yaml'
-    ])
+    mockedPickFiles.mockResolvedValueOnce(['/file.yaml', '/file2.yaml'])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: { val: [{ field1: 1, field2: 2 }] }, file: 'file2.yaml' },
-      { data: { val: [{ field1: 1 }] }, file: 'file2.yaml' }
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file.yaml' },
+      { data: { val: [{ field1: 1 }] }, file: '/file2.yaml' }
     ])
 
     const result = await pickPlotConfiguration('/')
@@ -217,8 +217,8 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedQuickPickValue
-      .mockResolvedValueOnce({ file: '/file.json', key: 'actual' })
-      .mockResolvedValueOnce({ file: '/file.json', key: 'prob' })
+      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
 
     const result = await pickPlotConfiguration('/')
 
@@ -245,8 +245,8 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: '/file.json', key: 'actual' } },
-        { label: 'prob', value: { file: '/file.json', key: 'prob' } }
+        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -258,7 +258,7 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'prob', value: { file: '/file.json', key: 'prob' } }
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -266,8 +266,8 @@ describe('pickPlotConfiguration', () => {
     )
     expect(result).toStrictEqual({
       template: 'simple',
-      x: { file: '/file.json', key: 'actual' },
-      y: { file: '/file.json', key: 'prob' }
+      x: { file: 'file.json', key: 'actual' },
+      y: { file: 'file.json', key: 'prob' }
     })
   })
 
@@ -279,8 +279,8 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedQuickPickValue
-      .mockResolvedValueOnce({ file: '/file.json', key: 'actual' })
-      .mockResolvedValueOnce({ file: '/file2.json', key: 'prob' })
+      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: 'file2.json', key: 'prob' })
 
     const result = await pickPlotConfiguration('/')
 
@@ -292,15 +292,15 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: '/file.json', key: 'actual' } },
-        { label: 'prob', value: { file: '/file.json', key: 'prob' } },
+        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: '/file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: '/file2.json', key: 'prob' } }
+        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -312,14 +312,14 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'prob', value: { file: '/file.json', key: 'prob' } },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: '/file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: '/file2.json', key: 'prob' } }
+        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -327,13 +327,13 @@ describe('pickPlotConfiguration', () => {
     )
     expect(result).toStrictEqual({
       template: 'simple',
-      x: { file: '/file.json', key: 'actual' },
-      y: { file: '/file2.json', key: 'prob' }
+      x: { file: 'file.json', key: 'actual' },
+      y: { file: 'file2.json', key: 'prob' }
     })
   })
 
   it('should return early if the user does not pick a template', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
@@ -347,7 +347,7 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should return early if the user does not pick a x field', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
@@ -361,7 +361,7 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should return early if the user does not pick a y field', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['file.json'])
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 190df91f53..56f865e477 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -23,13 +23,12 @@ type UnknownValue = Value | ValueTree
 
 type fileFields = { [file: string]: string[] }
 
-const pickDataFiles = (): Thenable<string[] | undefined> =>
+const pickDataFiles = (): Promise<string[] | undefined> =>
   pickFiles(Title.SELECT_PLOT_DATA, {
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
 
 const pickTemplateAndFields = async (
-  cwd: string,
   fields: fileFields
 ): Promise<PlotConfigData | undefined> => {
   const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
@@ -44,7 +43,7 @@ const pickTemplateAndFields = async (
     items.push(
       {
         kind: QuickPickItemKind.Separator,
-        label: relative(cwd, file),
+        label: file,
         value: undefined
       },
       ...keys.map(key => ({ label: key, value: { file, key } }))
@@ -81,7 +80,6 @@ const validateFileNames = (files: string[] | undefined) => {
   if (!files) {
     return []
   }
-
   const fileExts = [...new Set(files.map(file => getFileExtension(file)))]
 
   if (fileExts.length > 1) {
@@ -90,12 +88,12 @@ const validateFileNames = (files: string[] | undefined) => {
         fileExts
       )} extensions. Files must be of the same type.`
     )
-    return files
+    return []
   }
   return files
 }
 
-const getFieldsFromArr = (
+const getMetricInfoFromArr = (
   dataArr: UnknownValue[]
 ): { arrLength: number; fields: string[] } | undefined => {
   const firstArrVal: UnknownValue = dataArr[0]
@@ -112,7 +110,7 @@ const getFieldsFromArr = (
   return { arrLength: dataArr.length, fields: fieldObjKeys }
 }
 
-const getFieldsFromValue = (
+const getMetricInfoFromValue = (
   data: UnknownValue
 ): { arrLength: number; fields: string[] } | undefined => {
   const isArray = Array.isArray(data)
@@ -127,7 +125,7 @@ const getFieldsFromValue = (
     return
   }
 
-  return getFieldsFromArr(maybeFieldsObjArr)
+  return getMetricInfoFromArr(maybeFieldsObjArr)
 }
 
 const showNotEnoughKeysToast = (files: string[]) => {
@@ -148,7 +146,7 @@ const validateSingleDataFileFields = ({
   file: string
   data: UnknownValue
 }) => {
-  const { fields = [] } = getFieldsFromValue(data) || {}
+  const { fields = [] } = getMetricInfoFromValue(data) || {}
 
   if (fields.length < 2) {
     void showNotEnoughKeysToast([file])
@@ -166,14 +164,14 @@ const getFieldsFromDataFiles = (
   const keys: fileFields = {}
 
   for (const { file, data } of dataArr) {
-    const fileFields = getFieldsFromValue(data)
+    const metricInfo = getMetricInfoFromValue(data)
 
-    if (!fileFields) {
+    if (!metricInfo) {
       failedFiles.push(file)
       continue
     }
 
-    const { fields, arrLength } = fileFields
+    const { fields, arrLength } = metricInfo
 
     if (!keys[file]) {
       keys[file] = []
@@ -190,7 +188,7 @@ const validateMultiDataFileFields = (
 ) => {
   const { keys, failedFiles, filesArrLength } = getFieldsFromDataFiles(dataArr)
 
-  if (failedFiles) {
+  if (failedFiles.length > 0) {
     void showNotEnoughKeysToast(failedFiles)
     return
   }
@@ -204,13 +202,16 @@ const validateMultiDataFileFields = (
   return keys
 }
 
-const validateFilesData = async (files: string[]) => {
+const validateFilesData = async (cwd: string, files: string[]) => {
   const filesData = (await loadDataFiles(files)) as {
     data: UnknownValue
     file: string
   }[]
-
-  const failedFiles = filesData.filter(({ data }) => !!data)
+  const relativeFilesData = filesData.map(({ data, file }) => ({
+    data,
+    file: relative(cwd, file)
+  }))
+  const failedFiles = relativeFilesData.filter(({ data }) => !data)
 
   if (failedFiles.length > 0) {
     const files = joinList(failedFiles.map(({ file }) => file))
@@ -223,8 +224,8 @@ const validateFilesData = async (files: string[]) => {
   }
 
   return files.length === 1
-    ? validateSingleDataFileFields(filesData[0])
-    : validateMultiDataFileFields(filesData)
+    ? validateSingleDataFileFields(relativeFilesData[0])
+    : validateMultiDataFileFields(relativeFilesData)
 }
 
 export const pickPlotConfiguration = async (
@@ -237,13 +238,13 @@ export const pickPlotConfiguration = async (
     return
   }
 
-  const validFileFields = await validateFilesData(validFileNames)
+  const validFileFields = await validateFilesData(cwd, validFileNames)
 
   if (!validFileFields) {
     return
   }
 
-  const templateAndFields = await pickTemplateAndFields(cwd, validFileFields)
+  const templateAndFields = await pickTemplateAndFields(validFileFields)
 
   if (!templateAndFields) {
     return

From 528ec62626b949c089d4735ca1901280c3b26226 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 5 Oct 2023 11:31:58 -0500
Subject: [PATCH 15/35] Clean up and add more tests

---
 extension/src/pipeline/quickPick.test.ts | 61 +++++++++++++++++++++---
 extension/src/pipeline/quickPick.ts      | 57 ++++++++++++----------
 2 files changed, 86 insertions(+), 32 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 66d292761e..8f9c3ab0e6 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -81,7 +81,7 @@ describe('pickPlotConfiguration', () => {
     )
   })
 
-  it('should show a toast message if the file or files fail to parse', async () => {
+  it('should show a toast message if a single file to parse', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/data.csv'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: undefined, file: '/data.csv' }
@@ -92,7 +92,30 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'Failed to parse data.csv. Do the files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'Failed to parse data.csv. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+    )
+  })
+
+  it('should show a toast message if a multiple files to parse', async () => {
+    mockedPickFiles.mockResolvedValueOnce([
+      '/data.csv',
+      '/nested/values.csv',
+      '/results.csv',
+      '/props.csv'
+    ])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: undefined, file: '/data.csv' },
+      { data: undefined, file: '/nested/values.csv' },
+      { data: [{ field1: 'only one field' }], file: '/results.csv' },
+      { data: undefined, file: '/props.csv' }
+    ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+    expect(mockedShowError).toHaveBeenCalledWith(
+      'Failed to parse data.csv, nested/values.csv, and props.csv. Do the files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
@@ -154,9 +177,33 @@ describe('pickPlotConfiguration', () => {
       '/file3.yaml'
     ])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file2.yaml' },
-      { data: [], file: '/file3.yaml' },
-      { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file2.yaml' }
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file.yaml' },
+      { data: [], file: '/file2.yaml' },
+      { data: { val: 2 }, file: '/file3.yaml' }
+    ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+    expect(mockedShowError).toHaveBeenCalledWith(
+      'file2.yaml and file3.yaml do not contain any keys (columns) for plot generation. Do the files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+    )
+  })
+
+  it("should show a toast message if the multiple chosen files' object arrays are not the same length", async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.yaml', '/file2.yaml'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file.yaml' },
+      {
+        data: {
+          val: [
+            { field1: 1, field2: 2 },
+            { field1: 1, field2: 2 }
+          ]
+        },
+        file: '/file3.yaml'
+      }
     ])
 
     const result = await pickPlotConfiguration('/')
@@ -164,7 +211,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'file3.yaml does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'All files must have the same array (row) length. Do the files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
@@ -197,7 +244,7 @@ describe('pickPlotConfiguration', () => {
     }
   })
 
-  it('should parse fields from multiple valid data files (if atleast two fields are found total)', async () => {
+  it('should parse fields from multiple valid data files', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.yaml', '/file2.yaml'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: { val: [{ field1: 1, field2: 2 }] }, file: '/file.yaml' },
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 56f865e477..d538ac514c 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -73,7 +73,7 @@ const joinList = (items: string[]) => {
     return items.join(' and ')
   }
 
-  return `${[items.slice(0, -1)].join(', ')}, and ${items[items.length - 1]}`
+  return `${items.slice(0, -1).join(', ')}, and ${items[items.length - 1]}`
 }
 
 const validateFileNames = (files: string[] | undefined) => {
@@ -128,16 +128,10 @@ const getMetricInfoFromValue = (
   return getMetricInfoFromArr(maybeFieldsObjArr)
 }
 
-const showNotEnoughKeysToast = (files: string[]) => {
-  const isSingle = files.length === 1
-  return Toast.showError(
-    `${joinList(files)} ${
-      isSingle ? 'does' : 'do'
-    } not contain enough keys (columns) to generate a plot. Does the ${
-      isSingle ? 'file' : 'files'
-    } follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?`
-  )
-}
+const getFileGuidelinesQuestion = (isSingleFile: boolean) =>
+  `${
+    isSingleFile ? 'Does the file' : 'Do the files'
+  } follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?`
 
 const validateSingleDataFileFields = ({
   file,
@@ -149,17 +143,21 @@ const validateSingleDataFileFields = ({
   const { fields = [] } = getMetricInfoFromValue(data) || {}
 
   if (fields.length < 2) {
-    void showNotEnoughKeysToast([file])
+    void Toast.showError(
+      `${file} does not contain enough keys (columns) to generate a plot. ${getFileGuidelinesQuestion(
+        true
+      )}`
+    )
     return
   }
 
   return { [file]: fields }
 }
 
-const getFieldsFromDataFiles = (
+const getMetricInfoFromDataFiles = (
   dataArr: { data: UnknownValue; file: string }[]
 ) => {
-  const failedFiles: string[] = []
+  const invalidFiles: string[] = []
   const filesArrLength: Set<number> = new Set()
   const keys: fileFields = {}
 
@@ -167,7 +165,7 @@ const getFieldsFromDataFiles = (
     const metricInfo = getMetricInfoFromValue(data)
 
     if (!metricInfo) {
-      failedFiles.push(file)
+      invalidFiles.push(file)
       continue
     }
 
@@ -180,22 +178,31 @@ const getFieldsFromDataFiles = (
     keys[file].push(...fields)
   }
 
-  return { failedFiles, filesArrLength, keys }
+  return { filesArrLength, invalidFiles, keys }
 }
 
 const validateMultiDataFileFields = (
   dataArr: { data: UnknownValue; file: string }[]
 ) => {
-  const { keys, failedFiles, filesArrLength } = getFieldsFromDataFiles(dataArr)
+  const { keys, invalidFiles, filesArrLength } =
+    getMetricInfoFromDataFiles(dataArr)
 
-  if (failedFiles.length > 0) {
-    void showNotEnoughKeysToast(failedFiles)
+  if (invalidFiles.length > 0) {
+    void Toast.showError(
+      `${joinList(
+        invalidFiles
+      )} do not contain any keys (columns) for plot generation. ${getFileGuidelinesQuestion(
+        invalidFiles.length === 1
+      )}`
+    )
     return
   }
 
   if (filesArrLength.size > 1) {
     void Toast.showError(
-      'The files must have the same array (or row in tsv/csv) length.'
+      `All files must have the same array (row) length. ${getFileGuidelinesQuestion(
+        false
+      )}`
     )
   }
 
@@ -211,13 +218,13 @@ const validateFilesData = async (cwd: string, files: string[]) => {
     data,
     file: relative(cwd, file)
   }))
-  const failedFiles = relativeFilesData.filter(({ data }) => !data)
+  const invalidFilesData = relativeFilesData.filter(({ data }) => !data)
 
-  if (failedFiles.length > 0) {
-    const files = joinList(failedFiles.map(({ file }) => file))
+  if (invalidFilesData.length > 0) {
+    const invalidFiles = joinList(invalidFilesData.map(({ file }) => file))
     void Toast.showError(
-      `Failed to parse ${files}. ${
-        files.length === 1 ? 'Does the file' : 'Do the files'
+      `Failed to parse ${invalidFiles}. ${
+        invalidFilesData.length === 1 ? 'Does the file' : 'Do the files'
       } contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?`
     )
     return false

From 0d287a2f27acc87ba7d11afae107b346134d1b98 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 6 Oct 2023 12:15:40 -0500
Subject: [PATCH 16/35] Create more solid split between single and multi files

---
 extension/src/pipeline/quickPick.test.ts |  45 ++++------
 extension/src/pipeline/quickPick.ts      | 106 ++++++++---------------
 2 files changed, 57 insertions(+), 94 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 8f9c3ab0e6..76e83a9a64 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -72,42 +72,35 @@ describe('pickPlotConfiguration', () => {
   it('should show a toast message if the files are not the same data type', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file.csv'])
 
-    const result = await pickPlotConfiguration('/')
+    const twoExtsResult = await pickPlotConfiguration('/')
 
-    expect(result).toStrictEqual(undefined)
+    expect(twoExtsResult).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
       'Found files with .json and .csv extensions. Files must be of the same type.'
     )
-  })
 
-  it('should show a toast message if a single file to parse', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['/data.csv'])
-    mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: undefined, file: '/data.csv' }
+    mockedPickFiles.mockResolvedValueOnce([
+      '/file.json',
+      '/file.csv',
+      '/file.tsv'
     ])
 
-    const result = await pickPlotConfiguration('/')
+    const threeExtsResult = await pickPlotConfiguration('/')
 
-    expect(result).toStrictEqual(undefined)
-    expect(mockedShowError).toHaveBeenCalledTimes(1)
-    expect(mockedShowError).toHaveBeenCalledWith(
-      'Failed to parse data.csv. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+    expect(threeExtsResult).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(2)
+    expect(mockedShowError).toHaveBeenNthCalledWith(
+      2,
+      'Found files with .json, .csv, and .tsv extensions. Files must be of the same type.'
     )
   })
 
-  it('should show a toast message if a multiple files to parse', async () => {
-    mockedPickFiles.mockResolvedValueOnce([
-      '/data.csv',
-      '/nested/values.csv',
-      '/results.csv',
-      '/props.csv'
-    ])
+  it('should show a toast message if a file to parse', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/data.csv'])
     mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: undefined, file: '/data.csv' },
-      { data: undefined, file: '/nested/values.csv' },
       { data: [{ field1: 'only one field' }], file: '/results.csv' },
-      { data: undefined, file: '/props.csv' }
+      { data: undefined, file: '/data.csv' }
     ])
 
     const result = await pickPlotConfiguration('/')
@@ -115,7 +108,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'Failed to parse data.csv, nested/values.csv, and props.csv. Do the files contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'Failed to parse data.csv. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
@@ -165,7 +158,7 @@ describe('pickPlotConfiguration', () => {
       expect(result).toStrictEqual(undefined)
       expect(mockedShowError).toHaveBeenCalledTimes(1 + ind)
       expect(mockedShowError).toHaveBeenCalledWith(
-        'file.yaml does not contain enough keys (columns) to generate a plot. Does the file follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+        'file.yaml does not contain enough keys (columns) to generate a plot. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
       )
     }
   })
@@ -187,7 +180,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'file2.yaml and file3.yaml do not contain any keys (columns) for plot generation. Do the files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'file2.yaml does not contain any keys (columns) to generate a plot. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
@@ -211,7 +204,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'All files must have the same array (row) length. Do the files follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'All files must have the same array (list) length. See [examples](https://dvc.org/doc/command-reference/plots/show#sourcing-x-and-y-from-different-files) of multiple files being used in DVC plots.'
     )
   })
 
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index d538ac514c..ec1ab5983e 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -21,7 +21,7 @@ export type PlotConfigData = {
 
 type UnknownValue = Value | ValueTree
 
-type fileFields = { [file: string]: string[] }
+type FileFields = { file: string; fields: string[] }[]
 
 const pickDataFiles = (): Promise<string[] | undefined> =>
   pickFiles(Title.SELECT_PLOT_DATA, {
@@ -29,7 +29,7 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
   })
 
 const pickTemplateAndFields = async (
-  fields: fileFields
+  fileFields: FileFields
 ): Promise<PlotConfigData | undefined> => {
   const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
 
@@ -39,14 +39,14 @@ const pickTemplateAndFields = async (
 
   const items = []
 
-  for (const [file, keys] of Object.entries(fields)) {
+  for (const { file, fields } of fileFields) {
     items.push(
       {
         kind: QuickPickItemKind.Separator,
         label: file,
         value: undefined
       },
-      ...keys.map(key => ({ label: key, value: { file, key } }))
+      ...fields.map(key => ({ label: key, value: { file, key } }))
     )
   }
 
@@ -128,82 +128,54 @@ const getMetricInfoFromValue = (
   return getMetricInfoFromArr(maybeFieldsObjArr)
 }
 
-const getFileGuidelinesQuestion = (isSingleFile: boolean) =>
-  `${
-    isSingleFile ? 'Does the file' : 'Do the files'
-  } follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?`
+const getDvcGuidelinesQuestion = () =>
+  'Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
 
-const validateSingleDataFileFields = ({
+const validateSingleFileData = ({
   file,
   data
 }: {
-  file: string
   data: UnknownValue
+  file: string
 }) => {
   const { fields = [] } = getMetricInfoFromValue(data) || {}
 
   if (fields.length < 2) {
     void Toast.showError(
-      `${file} does not contain enough keys (columns) to generate a plot. ${getFileGuidelinesQuestion(
-        true
-      )}`
+      `${file} does not contain enough keys (columns) to generate a plot. ${getDvcGuidelinesQuestion()}`
     )
     return
   }
 
-  return { [file]: fields }
+  return [{ fields, file }]
 }
 
-const getMetricInfoFromDataFiles = (
-  dataArr: { data: UnknownValue; file: string }[]
+const validateMultiFilesData = (
+  filesData: { data: UnknownValue; file: string }[]
 ) => {
-  const invalidFiles: string[] = []
   const filesArrLength: Set<number> = new Set()
-  const keys: fileFields = {}
+  const keys: FileFields = []
 
-  for (const { file, data } of dataArr) {
+  for (const { file, data } of filesData) {
     const metricInfo = getMetricInfoFromValue(data)
-
     if (!metricInfo) {
-      invalidFiles.push(file)
-      continue
+      void Toast.showError(
+        `${file} does not contain any keys (columns) to generate a plot. ${getDvcGuidelinesQuestion()}`
+      )
+      return
     }
 
-    const { fields, arrLength } = metricInfo
-
-    if (!keys[file]) {
-      keys[file] = []
-    }
+    const { arrLength, fields } = metricInfo
+    keys.push({ fields, file })
     filesArrLength.add(arrLength)
-    keys[file].push(...fields)
-  }
-
-  return { filesArrLength, invalidFiles, keys }
-}
-
-const validateMultiDataFileFields = (
-  dataArr: { data: UnknownValue; file: string }[]
-) => {
-  const { keys, invalidFiles, filesArrLength } =
-    getMetricInfoFromDataFiles(dataArr)
-
-  if (invalidFiles.length > 0) {
-    void Toast.showError(
-      `${joinList(
-        invalidFiles
-      )} do not contain any keys (columns) for plot generation. ${getFileGuidelinesQuestion(
-        invalidFiles.length === 1
-      )}`
-    )
-    return
   }
 
   if (filesArrLength.size > 1) {
     void Toast.showError(
-      `All files must have the same array (row) length. ${getFileGuidelinesQuestion(
-        false
-      )}`
+      'All files must have the same array (list) length. See [examples](https://dvc.org/doc/command-reference/plots/show#sourcing-x-and-y-from-different-files) of multiple files being used in DVC plots.'
     )
+
+    return
   }
 
   return keys
@@ -214,25 +186,23 @@ const validateFilesData = async (cwd: string, files: string[]) => {
     data: UnknownValue
     file: string
   }[]
-  const relativeFilesData = filesData.map(({ data, file }) => ({
-    data,
-    file: relative(cwd, file)
-  }))
-  const invalidFilesData = relativeFilesData.filter(({ data }) => !data)
-
-  if (invalidFilesData.length > 0) {
-    const invalidFiles = joinList(invalidFilesData.map(({ file }) => file))
-    void Toast.showError(
-      `Failed to parse ${invalidFiles}. ${
-        invalidFilesData.length === 1 ? 'Does the file' : 'Do the files'
-      } contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?`
-    )
-    return false
+  const relativeFilesData = []
+
+  for (const { file, data } of filesData) {
+    const relativeFile = relative(cwd, file)
+    if (!data) {
+      void Toast.showError(
+        `Failed to parse ${relativeFile}. ${getDvcGuidelinesQuestion()}`
+      )
+      return
+    }
+
+    relativeFilesData.push({ data, file: relativeFile })
   }
 
-  return files.length === 1
-    ? validateSingleDataFileFields(relativeFilesData[0])
-    : validateMultiDataFileFields(relativeFilesData)
+  return relativeFilesData.length === 1
+    ? validateSingleFileData(relativeFilesData[0])
+    : validateMultiFilesData(relativeFilesData)
 }
 
 export const pickPlotConfiguration = async (

From bb688f9c454651686f7681ee0b51299e55e227f9 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 6 Oct 2023 18:17:29 -0500
Subject: [PATCH 17/35] Change function to constant

---
 extension/src/pipeline/quickPick.ts | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index ec1ab5983e..78149b7233 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -128,7 +128,7 @@ const getMetricInfoFromValue = (
   return getMetricInfoFromArr(maybeFieldsObjArr)
 }
 
-const getDvcGuidelinesQuestion = () =>
+const dvcPlotGuidelinesQuestion =
   'Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
 
 const validateSingleFileData = ({
@@ -142,7 +142,7 @@ const validateSingleFileData = ({
 
   if (fields.length < 2) {
     void Toast.showError(
-      `${file} does not contain enough keys (columns) to generate a plot. ${getDvcGuidelinesQuestion()}`
+      `${file} does not contain enough keys (columns) to generate a plot. ${dvcPlotGuidelinesQuestion}`
     )
     return
   }
@@ -160,7 +160,7 @@ const validateMultiFilesData = (
     const metricInfo = getMetricInfoFromValue(data)
     if (!metricInfo) {
       void Toast.showError(
-        `${file} does not contain any keys (columns) to generate a plot. ${getDvcGuidelinesQuestion()}`
+        `${file} does not contain any keys (columns) to generate a plot. ${dvcPlotGuidelinesQuestion}`
       )
       return
     }
@@ -192,7 +192,7 @@ const validateFilesData = async (cwd: string, files: string[]) => {
     const relativeFile = relative(cwd, file)
     if (!data) {
       void Toast.showError(
-        `Failed to parse ${relativeFile}. ${getDvcGuidelinesQuestion()}`
+        `Failed to parse ${relativeFile}. ${dvcPlotGuidelinesQuestion}`
       )
       return
     }

From 86556fcdf93e69ed888ffc94ed4fbcb847727ac2 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Fri, 6 Oct 2023 18:43:12 -0500
Subject: [PATCH 18/35] Stop returning array in `validateFileNames`

---
 extension/src/pipeline/quickPick.ts | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 78149b7233..1a16899bb1 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -78,7 +78,7 @@ const joinList = (items: string[]) => {
 
 const validateFileNames = (files: string[] | undefined) => {
   if (!files) {
-    return []
+    return
   }
   const fileExts = [...new Set(files.map(file => getFileExtension(file)))]
 
@@ -88,7 +88,7 @@ const validateFileNames = (files: string[] | undefined) => {
         fileExts
       )} extensions. Files must be of the same type.`
     )
-    return []
+    return
   }
   return files
 }
@@ -211,7 +211,7 @@ export const pickPlotConfiguration = async (
   const files = await pickDataFiles()
   const validFileNames = validateFileNames(files)
 
-  if (validFileNames.length === 0) {
+  if (!validFileNames) {
     return
   }
 

From d850b225ef6edb8154ff2becac3132ade4efb4a1 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 9 Oct 2023 08:21:06 -0500
Subject: [PATCH 19/35] Resolve review comments * fix typo * use same keys
 message for multi and single files

---
 extension/src/pipeline/quickPick.test.ts |  4 ++--
 extension/src/pipeline/quickPick.ts      | 13 +++++++------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 76e83a9a64..750075cedc 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -96,7 +96,7 @@ describe('pickPlotConfiguration', () => {
     )
   })
 
-  it('should show a toast message if a file to parse', async () => {
+  it('should show a toast message if parsing a file failed', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/data.csv'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: [{ field1: 'only one field' }], file: '/results.csv' },
@@ -180,7 +180,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
     expect(mockedShowError).toHaveBeenCalledTimes(1)
     expect(mockedShowError).toHaveBeenCalledWith(
-      'file2.yaml does not contain any keys (columns) to generate a plot. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
+      'file2.yaml does not contain enough keys (columns) to generate a plot. Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
     )
   })
 
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 1a16899bb1..63db921901 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -131,6 +131,11 @@ const getMetricInfoFromValue = (
 const dvcPlotGuidelinesQuestion =
   'Does the file contain data and follow the DVC plot guidelines for [JSON/YAML](https://dvc.org/doc/command-reference/plots/show#example-hierarchical-data) or [CSV/TSV](https://dvc.org/doc/command-reference/plots/show#example-tabular-data) files?'
 
+const showNotEnoughKeysToast = (file: string) =>
+  Toast.showError(
+    `${file} does not contain enough keys (columns) to generate a plot. ${dvcPlotGuidelinesQuestion}`
+  )
+
 const validateSingleFileData = ({
   file,
   data
@@ -141,9 +146,7 @@ const validateSingleFileData = ({
   const { fields = [] } = getMetricInfoFromValue(data) || {}
 
   if (fields.length < 2) {
-    void Toast.showError(
-      `${file} does not contain enough keys (columns) to generate a plot. ${dvcPlotGuidelinesQuestion}`
-    )
+    void showNotEnoughKeysToast(file)
     return
   }
 
@@ -159,9 +162,7 @@ const validateMultiFilesData = (
   for (const { file, data } of filesData) {
     const metricInfo = getMetricInfoFromValue(data)
     if (!metricInfo) {
-      void Toast.showError(
-        `${file} does not contain any keys (columns) to generate a plot. ${dvcPlotGuidelinesQuestion}`
-      )
+      void showNotEnoughKeysToast(file)
       return
     }
 

From cf57f3e7828f3e440324351f27ee3fe2de3070fc Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 9 Oct 2023 09:30:06 -0500
Subject: [PATCH 20/35] Add Title Option to Plot Wizard

---
 extension/src/fileSystem/index.test.ts        |  7 ++-
 extension/src/fileSystem/index.ts             |  5 +-
 extension/src/pipeline/quickPick.test.ts      | 29 ++++++++++-
 extension/src/pipeline/quickPick.ts           | 48 ++++++++++++++-----
 .../src/test/suite/pipeline/index.test.ts     |  1 +
 extension/src/vscode/title.ts                 |  1 +
 6 files changed, 75 insertions(+), 16 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 4fc43a5145..9a22fef5e4 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -608,7 +608,7 @@ describe('addPlotToDvcYamlFile', () => {
     '        eval/prc/test.json: precision'
   ]
   const mockNewPlotLines = [
-    '  - simple_plot:',
+    '  - Simple Plot:',
     '      template: simple',
     '      x: epochs',
     '      y:',
@@ -624,6 +624,7 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
+      title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
       y: { file: 'data.json', key: 'accuracy' }
     })
@@ -652,6 +653,7 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
+      title: 'simple_plot',
       x: { file: 'data.json', key: 'epochs' },
       y: { file: 'acc.json', key: 'accuracy' }
     })
@@ -670,6 +672,7 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
+      title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
       y: { file: 'data.json', key: 'accuracy' }
     })
@@ -692,6 +695,7 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
+      title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
       y: { file: 'data.json', key: 'accuracy' }
     })
@@ -729,6 +733,7 @@ describe('addPlotToDvcYamlFile', () => {
 
     addPlotToDvcYamlFile('/', {
       template: 'simple',
+      title: 'simple_plot',
       x: { file: 'data.json', key: 'epochs' },
       y: { file: 'data.json', key: 'accuracy' }
     })
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index c2ce48ce1b..a51c626abf 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -215,10 +215,9 @@ const loadYamlAsDoc = (
 }
 
 const getPlotYamlObj = (plot: PlotConfigData) => {
-  const { x, y, template } = plot
-  const plotName = `${template}_plot`
+  const { x, y, template, title } = plot
   return {
-    [plotName]: {
+    [title]: {
       template,
       x: x.file === y.file ? x.key : { [x.file]: x.key },
       y: { [y.file]: y.key }
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 750075cedc..b5a72dcc5c 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -1,5 +1,6 @@
 import { QuickPickItemKind } from 'vscode'
 import { pickPlotConfiguration } from './quickPick'
+import { getInput } from '../vscode/inputBox'
 import { pickFiles } from '../vscode/resourcePicker'
 import { quickPickOne, quickPickValue } from '../vscode/quickPick'
 import { getFileExtension, loadDataFiles } from '../fileSystem'
@@ -9,6 +10,7 @@ import { Toast } from '../vscode/toast'
 const mockedPickFiles = jest.mocked(pickFiles)
 const mockedLoadDataFiles = jest.mocked(loadDataFiles)
 const mockedGetFileExt = jest.mocked(getFileExtension)
+const mockedGetInput = jest.mocked(getInput)
 const mockedQuickPickOne = jest.mocked(quickPickOne)
 const mockedQuickPickValue = jest.mocked(quickPickValue)
 const mockedToast = jest.mocked(Toast)
@@ -16,6 +18,7 @@ const mockedShowError = jest.fn()
 mockedToast.showError = mockedShowError
 
 jest.mock('../fileSystem')
+jest.mock('../vscode/inputBox')
 jest.mock('../vscode/resourcePicker')
 jest.mock('../vscode/quickPick')
 
@@ -250,18 +253,23 @@ describe('pickPlotConfiguration', () => {
     expect(mockedShowError).not.toHaveBeenCalled()
   })
 
-  it('should let the user pick a template, x field, and y field', async () => {
+  it('should let the user pick a template, title, x field, and y field', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: '/file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('Simple Plot')
     mockedQuickPickValue
       .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
       .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
 
     const result = await pickPlotConfiguration('/')
 
+    expect(mockedGetInput).toHaveBeenCalledWith(
+      Title.ENTER_PLOT_TITLE,
+      'simple_plot'
+    )
     expect(mockedQuickPickOne).toHaveBeenNthCalledWith(
       1,
       [
@@ -306,6 +314,7 @@ describe('pickPlotConfiguration', () => {
     )
     expect(result).toStrictEqual({
       template: 'simple',
+      title: 'Simple Plot',
       x: { file: 'file.json', key: 'actual' },
       y: { file: 'file.json', key: 'prob' }
     })
@@ -318,6 +327,7 @@ describe('pickPlotConfiguration', () => {
       { data: mockValidData, file: '/file2.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickValue
       .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
       .mockResolvedValueOnce({ file: 'file2.json', key: 'prob' })
@@ -367,6 +377,7 @@ describe('pickPlotConfiguration', () => {
     )
     expect(result).toStrictEqual({
       template: 'simple',
+      title: 'simple_plot',
       x: { file: 'file.json', key: 'actual' },
       y: { file: 'file2.json', key: 'prob' }
     })
@@ -386,12 +397,27 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
   })
 
+  it('should return early if the user does not pick a title', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('linear')
+    mockedGetInput.mockResolvedValueOnce(undefined)
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedGetInput).toHaveBeenCalledTimes(1)
+    expect(result).toStrictEqual(undefined)
+  })
+
   it('should return early if the user does not pick a x field', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickValue.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
@@ -406,6 +432,7 @@ describe('pickPlotConfiguration', () => {
       { data: mockValidData, file: 'file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickValue
       .mockResolvedValueOnce('actual')
       .mockResolvedValueOnce(undefined)
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 63db921901..c4f58b88c3 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -12,10 +12,12 @@ import { quickPickOne, quickPickValue } from '../vscode/quickPick'
 import { pickFiles } from '../vscode/resourcePicker'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
+import { getInput } from '../vscode/inputBox'
 
 export type PlotConfigData = {
   x: { file: string; key: string }
   template: string
+  title: string
   y: { file: string; key: string }
 }
 
@@ -28,15 +30,27 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
 
-const pickTemplateAndFields = async (
-  fileFields: FileFields
-): Promise<PlotConfigData | undefined> => {
+const pickTemplateAndTitle = async (): Promise<
+  Omit<PlotConfigData, 'x' | 'y'> | undefined
+> => {
   const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
 
   if (!template) {
     return
   }
 
+  const title = await getInput(Title.ENTER_PLOT_TITLE, `${template}_plot`)
+
+  if (!title) {
+    return
+  }
+
+  return { template, title }
+}
+
+const pickFields = async (
+  fileFields: FileFields
+): Promise<Omit<PlotConfigData, 'title' | 'template'> | undefined> => {
   const items = []
 
   for (const { file, fields } of fileFields) {
@@ -65,7 +79,25 @@ const pickTemplateAndFields = async (
     return
   }
 
-  return { template, x, y }
+  return { x, y }
+}
+
+const pickPlotConfigData = async (
+  fileFields: FileFields
+): Promise<PlotConfigData | undefined> => {
+  const templateAndTitle = await pickTemplateAndTitle()
+
+  if (!templateAndTitle) {
+    return
+  }
+
+  const fields = await pickFields(fileFields)
+
+  if (!fields) {
+    return
+  }
+
+  return { ...templateAndTitle, ...fields }
 }
 
 const joinList = (items: string[]) => {
@@ -222,11 +254,5 @@ export const pickPlotConfiguration = async (
     return
   }
 
-  const templateAndFields = await pickTemplateAndFields(validFileFields)
-
-  if (!templateAndFields) {
-    return
-  }
-
-  return { ...templateAndFields }
+  return pickPlotConfigData(validFileFields)
 }
diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts
index 9ba5a92915..74a6cf0c99 100644
--- a/extension/src/test/suite/pipeline/index.test.ts
+++ b/extension/src/test/suite/pipeline/index.test.ts
@@ -311,6 +311,7 @@ suite('Pipeline Test Suite', () => {
 
     mockPickPlotConfiguration.onSecondCall().resolves({
       template: 'simple',
+      title: 'Great Plot Name',
       x: { file: 'results.json', key: 'step' },
       y: { file: 'results.json', key: 'acc' }
     })
diff --git a/extension/src/vscode/title.ts b/extension/src/vscode/title.ts
index 61a6f34de2..9b1fdad5b4 100644
--- a/extension/src/vscode/title.ts
+++ b/extension/src/vscode/title.ts
@@ -10,6 +10,7 @@ export enum Title {
   ENTER_REMOTE_NAME = 'Enter a Name for the Remote',
   ENTER_REMOTE_URL = 'Enter the URL for the Remote',
   ENTER_PATH_OR_CHOOSE_FILE = 'Enter the path to your training script or select it',
+  ENTER_PLOT_TITLE = 'Enter a Name for Your Plot',
   ENTER_STUDIO_USERNAME = 'Enter your Studio username',
   ENTER_STUDIO_TOKEN = 'Enter your Studio access token',
   ENTER_STAGE_NAME = 'Enter a Name for the Main Stage of your Pipeline',

From 89fc205f91838081d428327f18cbe92b82196774 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 9 Oct 2023 10:35:56 -0500
Subject: [PATCH 21/35] Add Multiple Y Field Selection in Plot Wizard

---
 extension/src/fileSystem/index.test.ts        |  10 +-
 extension/src/fileSystem/index.ts             |  14 ++-
 extension/src/pipeline/quickPick.test.ts      | 116 ++++++++++++++----
 extension/src/pipeline/quickPick.ts           |  24 +++-
 .../src/test/suite/pipeline/index.test.ts     |   2 +-
 5 files changed, 132 insertions(+), 34 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 9a22fef5e4..a1fe4ff820 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -626,7 +626,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { file: 'data.json', key: 'accuracy' }
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -655,7 +655,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'simple_plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { file: 'acc.json', key: 'accuracy' }
+      y: { 'acc.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -674,7 +674,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { file: 'data.json', key: 'accuracy' }
+      y: { 'data.json': ['accuracy'] }
     })
 
     mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
@@ -697,7 +697,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { file: 'data.json', key: 'accuracy' }
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -735,7 +735,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'simple_plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { file: 'data.json', key: 'accuracy' }
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index a51c626abf..9e4867e951 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -216,11 +216,21 @@ const loadYamlAsDoc = (
 
 const getPlotYamlObj = (plot: PlotConfigData) => {
   const { x, y, template, title } = plot
+
+  const yFiles = Object.keys(y)
+  const oneFileUsed = yFiles.length === 1 && yFiles[0] === x.file
+
+  const formattedY: { [file: string]: string | string[] } = {}
+
+  for (const [file, keys] of Object.entries(y)) {
+    formattedY[file] = keys.length === 1 ? keys[0] : keys
+  }
+
   return {
     [title]: {
       template,
-      x: x.file === y.file ? x.key : { [x.file]: x.key },
-      y: { [y.file]: y.key }
+      x: oneFileUsed ? x.key : { [x.file]: x.key },
+      y: formattedY
     }
   }
 }
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index b5a72dcc5c..6d28b7f72f 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -2,7 +2,11 @@ import { QuickPickItemKind } from 'vscode'
 import { pickPlotConfiguration } from './quickPick'
 import { getInput } from '../vscode/inputBox'
 import { pickFiles } from '../vscode/resourcePicker'
-import { quickPickOne, quickPickValue } from '../vscode/quickPick'
+import {
+  quickPickOne,
+  quickPickValue,
+  quickPickManyValues
+} from '../vscode/quickPick'
 import { getFileExtension, loadDataFiles } from '../fileSystem'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
@@ -13,6 +17,7 @@ const mockedGetFileExt = jest.mocked(getFileExtension)
 const mockedGetInput = jest.mocked(getInput)
 const mockedQuickPickOne = jest.mocked(quickPickOne)
 const mockedQuickPickValue = jest.mocked(quickPickValue)
+const mockedQuickPickValues = jest.mocked(quickPickManyValues)
 const mockedToast = jest.mocked(Toast)
 const mockedShowError = jest.fn()
 mockedToast.showError = mockedShowError
@@ -260,9 +265,13 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('Simple Plot')
-    mockedQuickPickValue
-      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
-      .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
+    mockedQuickPickValue.mockResolvedValueOnce({
+      file: 'file.json',
+      key: 'actual'
+    })
+    mockedQuickPickValues.mockResolvedValueOnce([
+      { file: 'file.json', key: 'prob' }
+    ])
 
     const result = await pickPlotConfiguration('/')
 
@@ -285,8 +294,7 @@ describe('pickPlotConfiguration', () => {
       ],
       'Pick a Plot Template'
     )
-    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
-      1,
+    expect(mockedQuickPickValue).toHaveBeenCalledWith(
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -298,8 +306,7 @@ describe('pickPlotConfiguration', () => {
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
-    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
-      2,
+    expect(mockedQuickPickValues).toHaveBeenCalledWith(
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -316,7 +323,7 @@ describe('pickPlotConfiguration', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'file.json', key: 'actual' },
-      y: { file: 'file.json', key: 'prob' }
+      y: { 'file.json': ['prob'] }
     })
   })
 
@@ -328,14 +335,17 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickValue
-      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
-      .mockResolvedValueOnce({ file: 'file2.json', key: 'prob' })
+    mockedQuickPickValue.mockResolvedValueOnce({
+      file: 'file.json',
+      key: 'actual'
+    })
+    mockedQuickPickValues.mockResolvedValueOnce([
+      { file: 'file2.json', key: 'prob' }
+    ])
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
-      1,
+    expect(mockedQuickPickValue).toHaveBeenCalledWith(
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -354,8 +364,73 @@ describe('pickPlotConfiguration', () => {
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
-    expect(mockedQuickPickValue).toHaveBeenNthCalledWith(
-      2,
+    expect(mockedQuickPickValues).toHaveBeenCalledWith(
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file2.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+      ],
+      {
+        title: Title.SELECT_PLOT_Y_METRIC
+      }
+    )
+    expect(result).toStrictEqual({
+      template: 'simple',
+      title: 'simple_plot',
+      x: { file: 'file.json', key: 'actual' },
+      y: { 'file2.json': ['prob'] }
+    })
+  })
+
+  it('should let the user pick multiple y fields', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file2.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: '/file.json' },
+      { data: mockValidData, file: '/file2.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickValue.mockResolvedValueOnce({
+      file: 'file.json',
+      key: 'actual'
+    })
+    mockedQuickPickValues.mockResolvedValueOnce([
+      { file: 'file2.json', key: 'prob' },
+      { file: 'file2.json', key: 'actual' }
+    ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickValue).toHaveBeenCalledWith(
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file2.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+      ],
+      { title: Title.SELECT_PLOT_X_METRIC }
+    )
+    expect(mockedQuickPickValues).toHaveBeenCalledWith(
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -379,7 +454,7 @@ describe('pickPlotConfiguration', () => {
       template: 'simple',
       title: 'simple_plot',
       x: { file: 'file.json', key: 'actual' },
-      y: { file: 'file2.json', key: 'prob' }
+      y: { 'file2.json': ['prob', 'actual'] }
     })
   })
 
@@ -433,13 +508,12 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickValue
-      .mockResolvedValueOnce('actual')
-      .mockResolvedValueOnce(undefined)
+    mockedQuickPickValue.mockResolvedValueOnce('actual')
+    mockedQuickPickValues.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValue).toHaveBeenCalledTimes(2)
+    expect(mockedQuickPickValues).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
   })
 })
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index c4f58b88c3..2e3224272c 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -8,7 +8,11 @@ import {
   isValueTree
 } from '../cli/dvc/contract'
 import { getFileExtension, loadDataFiles } from '../fileSystem'
-import { quickPickOne, quickPickValue } from '../vscode/quickPick'
+import {
+  quickPickManyValues,
+  quickPickOne,
+  quickPickValue
+} from '../vscode/quickPick'
 import { pickFiles } from '../vscode/resourcePicker'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
@@ -18,7 +22,7 @@ export type PlotConfigData = {
   x: { file: string; key: string }
   template: string
   title: string
-  y: { file: string; key: string }
+  y: { [file: string]: string[] }
 }
 
 type UnknownValue = Value | ValueTree
@@ -70,15 +74,25 @@ const pickFields = async (
     return
   }
 
-  const y = await quickPickValue(
+  const yValues = (await quickPickManyValues(
     items.filter(item => !isEqual(item.value, x)),
     { title: Title.SELECT_PLOT_Y_METRIC }
-  )
+  )) as { file: string; key: string }[] | undefined
 
-  if (!y) {
+  if (!yValues) {
     return
   }
 
+  const y: PlotConfigData['y'] = {}
+
+  for (const { file, key } of yValues) {
+    if (!y[file]) {
+      y[file] = []
+    }
+
+    y[file].push(key)
+  }
+
   return { x, y }
 }
 
diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts
index 74a6cf0c99..38ded7eab6 100644
--- a/extension/src/test/suite/pipeline/index.test.ts
+++ b/extension/src/test/suite/pipeline/index.test.ts
@@ -313,7 +313,7 @@ suite('Pipeline Test Suite', () => {
       template: 'simple',
       title: 'Great Plot Name',
       x: { file: 'results.json', key: 'step' },
-      y: { file: 'results.json', key: 'acc' }
+      y: { 'results.json': ['acc'] }
     })
 
     await pipeline.addDataSeriesPlot()

From 20f6ce70eb2398c1fc89d6a3d886fa70da2d62d1 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 9 Oct 2023 11:31:28 -0500
Subject: [PATCH 22/35] Only format y values once in `quickPick.ts`

---
 extension/src/fileSystem/index.test.ts   | 10 +++----
 extension/src/fileSystem/index.ts        |  8 +-----
 extension/src/pipeline/quickPick.test.ts | 25 ++++++++++++------
 extension/src/pipeline/quickPick.ts      | 33 +++++++++++++++---------
 4 files changed, 44 insertions(+), 32 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index a1fe4ff820..a434fe8e41 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -626,7 +626,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { 'data.json': ['accuracy'] }
+      y: { 'data.json': 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -655,7 +655,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'simple_plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { 'acc.json': ['accuracy'] }
+      y: { 'acc.json': 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -674,7 +674,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { 'data.json': ['accuracy'] }
+      y: { 'data.json': 'accuracy' }
     })
 
     mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
@@ -697,7 +697,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { 'data.json': ['accuracy'] }
+      y: { 'data.json': 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -735,7 +735,7 @@ describe('addPlotToDvcYamlFile', () => {
       template: 'simple',
       title: 'simple_plot',
       x: { file: 'data.json', key: 'epochs' },
-      y: { 'data.json': ['accuracy'] }
+      y: { 'data.json': 'accuracy' }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index 9e4867e951..a09122e6af 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -220,17 +220,11 @@ const getPlotYamlObj = (plot: PlotConfigData) => {
   const yFiles = Object.keys(y)
   const oneFileUsed = yFiles.length === 1 && yFiles[0] === x.file
 
-  const formattedY: { [file: string]: string | string[] } = {}
-
-  for (const [file, keys] of Object.entries(y)) {
-    formattedY[file] = keys.length === 1 ? keys[0] : keys
-  }
-
   return {
     [title]: {
       template,
       x: oneFileUsed ? x.key : { [x.file]: x.key },
-      y: formattedY
+      y
     }
   }
 }
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 6d28b7f72f..4332b35da8 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -323,7 +323,7 @@ describe('pickPlotConfiguration', () => {
       template: 'simple',
       title: 'Simple Plot',
       x: { file: 'file.json', key: 'actual' },
-      y: { 'file.json': ['prob'] }
+      y: { 'file.json': 'prob' }
     })
   })
 
@@ -331,7 +331,10 @@ describe('pickPlotConfiguration', () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file2.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: '/file.json' },
-      { data: mockValidData, file: '/file2.json' }
+      {
+        data: mockValidData,
+        file: '/file2.json'
+      }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
@@ -388,7 +391,7 @@ describe('pickPlotConfiguration', () => {
       template: 'simple',
       title: 'simple_plot',
       x: { file: 'file.json', key: 'actual' },
-      y: { 'file2.json': ['prob'] }
+      y: { 'file2.json': 'prob' }
     })
   })
 
@@ -396,7 +399,10 @@ describe('pickPlotConfiguration', () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file2.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: '/file.json' },
-      { data: mockValidData, file: '/file2.json' }
+      {
+        data: mockValidData.map((value, ind) => ({ ...value, step: ind })),
+        file: '/file2.json'
+      }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
@@ -406,7 +412,8 @@ describe('pickPlotConfiguration', () => {
     })
     mockedQuickPickValues.mockResolvedValueOnce([
       { file: 'file2.json', key: 'prob' },
-      { file: 'file2.json', key: 'actual' }
+      { file: 'file2.json', key: 'actual' },
+      { file: 'file2.json', key: 'step' }
     ])
 
     const result = await pickPlotConfiguration('/')
@@ -426,7 +433,8 @@ describe('pickPlotConfiguration', () => {
           value: undefined
         },
         { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } },
+        { label: 'step', value: { file: 'file2.json', key: 'step' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -444,7 +452,8 @@ describe('pickPlotConfiguration', () => {
           value: undefined
         },
         { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } },
+        { label: 'step', value: { file: 'file2.json', key: 'step' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -454,7 +463,7 @@ describe('pickPlotConfiguration', () => {
       template: 'simple',
       title: 'simple_plot',
       x: { file: 'file.json', key: 'actual' },
-      y: { 'file2.json': ['prob', 'actual'] }
+      y: { 'file2.json': ['prob', 'actual', 'step'] }
     })
   })
 
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 2e3224272c..c008c4c965 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -22,7 +22,7 @@ export type PlotConfigData = {
   x: { file: string; key: string }
   template: string
   title: string
-  y: { [file: string]: string[] }
+  y: { [file: string]: string[] | string }
 }
 
 type UnknownValue = Value | ValueTree
@@ -52,6 +52,25 @@ const pickTemplateAndTitle = async (): Promise<
   return { template, title }
 }
 
+const formatYQuickPickValues = (values: { file: string; key: string }[]) => {
+  const y: PlotConfigData['y'] = {}
+
+  for (const { file, key } of values) {
+    if (!y[file]) {
+      y[file] = key
+      continue
+    }
+
+    const prevFileValue = y[file]
+    y[file] = [
+      ...(typeof prevFileValue === 'string' ? [prevFileValue] : prevFileValue),
+      key
+    ]
+  }
+
+  return y
+}
+
 const pickFields = async (
   fileFields: FileFields
 ): Promise<Omit<PlotConfigData, 'title' | 'template'> | undefined> => {
@@ -83,17 +102,7 @@ const pickFields = async (
     return
   }
 
-  const y: PlotConfigData['y'] = {}
-
-  for (const { file, key } of yValues) {
-    if (!y[file]) {
-      y[file] = []
-    }
-
-    y[file].push(key)
-  }
-
-  return { x, y }
+  return { x, y: formatYQuickPickValues(yValues) }
 }
 
 const pickPlotConfigData = async (

From 9a185b0466e6077807b1121eea2a04c528f783f1 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Mon, 9 Oct 2023 11:45:34 -0500
Subject: [PATCH 23/35] Fix typo

---
 extension/src/test/suite/pipeline/index.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts
index 38ded7eab6..a2ed7a9d30 100644
--- a/extension/src/test/suite/pipeline/index.test.ts
+++ b/extension/src/test/suite/pipeline/index.test.ts
@@ -313,7 +313,7 @@ suite('Pipeline Test Suite', () => {
       template: 'simple',
       title: 'Great Plot Name',
       x: { file: 'results.json', key: 'step' },
-      y: { 'results.json': ['acc'] }
+      y: { 'results.json': 'acc' }
     })
 
     await pipeline.addDataSeriesPlot()

From 861fe520a22a69205a5ece9ec42dabdbaa209e66 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 10 Oct 2023 07:51:36 -0500
Subject: [PATCH 24/35] Update default title

---
 extension/src/pipeline/quickPick.test.ts | 41 +++++++++++++-----------
 extension/src/pipeline/quickPick.ts      | 33 +++++++------------
 2 files changed, 34 insertions(+), 40 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index b5a72dcc5c..a81b66ca72 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -253,23 +253,19 @@ describe('pickPlotConfiguration', () => {
     expect(mockedShowError).not.toHaveBeenCalled()
   })
 
-  it('should let the user pick a template, title, x field, and y field', async () => {
+  it('should let the user pick a template, x field, y field and title', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: '/file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
-    mockedGetInput.mockResolvedValueOnce('Simple Plot')
     mockedQuickPickValue
       .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
       .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
+    mockedGetInput.mockResolvedValueOnce('Simple Plot')
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedGetInput).toHaveBeenCalledWith(
-      Title.ENTER_PLOT_TITLE,
-      'simple_plot'
-    )
     expect(mockedQuickPickOne).toHaveBeenNthCalledWith(
       1,
       [
@@ -312,6 +308,10 @@ describe('pickPlotConfiguration', () => {
         title: Title.SELECT_PLOT_Y_METRIC
       }
     )
+    expect(mockedGetInput).toHaveBeenCalledWith(
+      Title.ENTER_PLOT_TITLE,
+      'actual vs prob'
+    )
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'Simple Plot',
@@ -397,49 +397,52 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
   })
 
-  it('should return early if the user does not pick a title', async () => {
+  it('should return early if the user does not pick a x field', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
-    mockedQuickPickOne.mockResolvedValueOnce('linear')
-    mockedGetInput.mockResolvedValueOnce(undefined)
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickValue.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedGetInput).toHaveBeenCalledTimes(1)
+    expect(mockedQuickPickValue).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
   })
 
-  it('should return early if the user does not pick a x field', async () => {
+  it('should return early if the user does not pick a y field', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickValue.mockResolvedValueOnce(undefined)
+    mockedQuickPickValue
+      .mockResolvedValueOnce('actual')
+      .mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValue).toHaveBeenCalledTimes(1)
+    expect(mockedQuickPickValue).toHaveBeenCalledTimes(2)
     expect(result).toStrictEqual(undefined)
   })
 
-  it('should return early if the user does not pick a y field', async () => {
+  it('should return early if the user does not pick a title', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
-    mockedQuickPickOne.mockResolvedValueOnce('simple')
-    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickOne.mockResolvedValueOnce('linear')
     mockedQuickPickValue
-      .mockResolvedValueOnce('actual')
-      .mockResolvedValueOnce(undefined)
+      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
+    mockedGetInput.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValue).toHaveBeenCalledTimes(2)
+    expect(mockedGetInput).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
   })
 })
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index c4f58b88c3..5ee6aa9ccc 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -30,24 +30,6 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
 
-const pickTemplateAndTitle = async (): Promise<
-  Omit<PlotConfigData, 'x' | 'y'> | undefined
-> => {
-  const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
-
-  if (!template) {
-    return
-  }
-
-  const title = await getInput(Title.ENTER_PLOT_TITLE, `${template}_plot`)
-
-  if (!title) {
-    return
-  }
-
-  return { template, title }
-}
-
 const pickFields = async (
   fileFields: FileFields
 ): Promise<Omit<PlotConfigData, 'title' | 'template'> | undefined> => {
@@ -85,9 +67,9 @@ const pickFields = async (
 const pickPlotConfigData = async (
   fileFields: FileFields
 ): Promise<PlotConfigData | undefined> => {
-  const templateAndTitle = await pickTemplateAndTitle()
+  const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
 
-  if (!templateAndTitle) {
+  if (!template) {
     return
   }
 
@@ -97,7 +79,16 @@ const pickPlotConfigData = async (
     return
   }
 
-  return { ...templateAndTitle, ...fields }
+  const title = await getInput(
+    Title.ENTER_PLOT_TITLE,
+    `${fields.x.key} vs ${fields.y.key}`
+  )
+
+  if (!title) {
+    return
+  }
+
+  return { template, title, ...fields }
 }
 
 const joinList = (items: string[]) => {

From 29657e490444f33132d44059b2c889a9cf630414 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 10 Oct 2023 08:26:33 -0500
Subject: [PATCH 25/35] Keep x the same type as y

---
 extension/src/fileSystem/index.test.ts   | 10 +++----
 extension/src/fileSystem/index.ts        |  5 ++--
 extension/src/pipeline/quickPick.test.ts |  6 ++--
 extension/src/pipeline/quickPick.ts      | 37 +++++++++++++++---------
 4 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index a434fe8e41..0cd03c57ed 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -625,7 +625,7 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { file: 'data.json', key: 'epochs' },
+      x: { 'data.json': 'epochs' },
       y: { 'data.json': 'accuracy' }
     })
 
@@ -654,7 +654,7 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'simple_plot',
-      x: { file: 'data.json', key: 'epochs' },
+      x: { 'data.json': 'epochs' },
       y: { 'acc.json': 'accuracy' }
     })
 
@@ -673,7 +673,7 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { file: 'data.json', key: 'epochs' },
+      x: { 'data.json': 'epochs' },
       y: { 'data.json': 'accuracy' }
     })
 
@@ -696,7 +696,7 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { file: 'data.json', key: 'epochs' },
+      x: { 'data.json': 'epochs' },
       y: { 'data.json': 'accuracy' }
     })
 
@@ -734,7 +734,7 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'simple_plot',
-      x: { file: 'data.json', key: 'epochs' },
+      x: { 'data.json': 'epochs' },
       y: { 'data.json': 'accuracy' }
     })
 
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index a09122e6af..72edb1ed4f 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -218,12 +218,13 @@ const getPlotYamlObj = (plot: PlotConfigData) => {
   const { x, y, template, title } = plot
 
   const yFiles = Object.keys(y)
-  const oneFileUsed = yFiles.length === 1 && yFiles[0] === x.file
+  const [xFile, xKey] = Object.entries(x)[0]
+  const oneFileUsed = yFiles.length === 1 && yFiles[0] === xFile
 
   return {
     [title]: {
       template,
-      x: oneFileUsed ? x.key : { [x.file]: x.key },
+      x: oneFileUsed ? xKey : x,
       y
     }
   }
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index c99cfb7e5f..913ad86a0d 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -322,7 +322,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'Simple Plot',
-      x: { file: 'file.json', key: 'actual' },
+      x: { 'file.json': 'actual' },
       y: { 'file.json': 'prob' }
     })
   })
@@ -390,7 +390,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'simple_plot',
-      x: { file: 'file.json', key: 'actual' },
+      x: { 'file.json': 'actual' },
       y: { 'file2.json': 'prob' }
     })
   })
@@ -462,7 +462,7 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'simple_plot',
-      x: { file: 'file.json', key: 'actual' },
+      x: { 'file.json': 'actual' },
       y: { 'file2.json': ['prob', 'actual', 'step'] }
     })
   })
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index dcd383de0a..88178ea5a7 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -19,7 +19,7 @@ import { Toast } from '../vscode/toast'
 import { getInput } from '../vscode/inputBox'
 
 export type PlotConfigData = {
-  x: { file: string; key: string }
+  x: { [file: string]: string }
   template: string
   title: string
   y: { [file: string]: string[] | string }
@@ -34,23 +34,25 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
 
-const formatYQuickPickValues = (values: { file: string; key: string }[]) => {
-  const y: PlotConfigData['y'] = {}
+const formatFieldQuickPickValues = (
+  values: { file: string; key: string }[]
+) => {
+  const formattedFields: PlotConfigData['y'] = {}
 
   for (const { file, key } of values) {
-    if (!y[file]) {
-      y[file] = key
+    if (!formattedFields[file]) {
+      formattedFields[file] = key
       continue
     }
 
-    const prevFileValue = y[file]
-    y[file] = [
+    const prevFileValue = formattedFields[file]
+    formattedFields[file] = [
       ...(typeof prevFileValue === 'string' ? [prevFileValue] : prevFileValue),
       key
     ]
   }
 
-  return y
+  return formattedFields
 }
 
 const pickFields = async (
@@ -58,6 +60,7 @@ const pickFields = async (
 ): Promise<
   | {
       fields: Omit<PlotConfigData, 'title' | 'template'>
+      firstXKey: string
       firstYKey: string
     }
   | undefined
@@ -75,14 +78,16 @@ const pickFields = async (
     )
   }
 
-  const x = await quickPickValue(items, { title: Title.SELECT_PLOT_X_METRIC })
+  const xValue = await quickPickValue(items, {
+    title: Title.SELECT_PLOT_X_METRIC
+  })
 
-  if (!x) {
+  if (!xValue) {
     return
   }
 
   const yValues = (await quickPickManyValues(
-    items.filter(item => !isEqual(item.value, x)),
+    items.filter(item => !isEqual(item.value, xValue)),
     { title: Title.SELECT_PLOT_Y_METRIC }
   )) as { file: string; key: string }[] | undefined
 
@@ -91,7 +96,11 @@ const pickFields = async (
   }
 
   return {
-    fields: { x, y: formatYQuickPickValues(yValues) },
+    fields: {
+      x: { [xValue.file]: xValue.key },
+      y: formatFieldQuickPickValues(yValues)
+    },
+    firstXKey: xValue.key,
     firstYKey: yValues[0].key
   }
 }
@@ -111,11 +120,11 @@ const pickPlotConfigData = async (
     return
   }
 
-  const { fields, firstYKey } = fieldsInfo
+  const { fields, firstYKey, firstXKey } = fieldsInfo
 
   const title = await getInput(
     Title.ENTER_PLOT_TITLE,
-    `${fields.x.key} vs ${firstYKey}`
+    `${firstXKey} vs ${firstYKey}`
   )
 
   if (!title) {

From 3b5bc205047f31f804cce180aab85300ef2a8962 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 10 Oct 2023 11:45:58 -0500
Subject: [PATCH 26/35] first iteration

---
 extension/src/fileSystem/index.ts        |   7 +-
 extension/src/pipeline/quickPick.test.ts | 274 +++++++++++++++++++----
 extension/src/pipeline/quickPick.ts      |  24 +-
 extension/src/vscode/quickPick.ts        |  40 ++++
 4 files changed, 285 insertions(+), 60 deletions(-)

diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index 72edb1ed4f..11f16862e6 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -218,13 +218,14 @@ const getPlotYamlObj = (plot: PlotConfigData) => {
   const { x, y, template, title } = plot
 
   const yFiles = Object.keys(y)
-  const [xFile, xKey] = Object.entries(x)[0]
-  const oneFileUsed = yFiles.length === 1 && yFiles[0] === xFile
+  const xFiles = Object.keys(x)
+  const oneFileUsed =
+    yFiles.length === 1 && xFiles.length === 1 && yFiles[0] === xFiles[0]
 
   return {
     [title]: {
       template,
-      x: oneFileUsed ? xKey : x,
+      x: oneFileUsed ? x[xFiles[0]] : x,
       y
     }
   }
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 913ad86a0d..1daeeeeade 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -5,7 +5,7 @@ import { pickFiles } from '../vscode/resourcePicker'
 import {
   quickPickOne,
   quickPickValue,
-  quickPickManyValues
+  quickPickUserOrderedValues
 } from '../vscode/quickPick'
 import { getFileExtension, loadDataFiles } from '../fileSystem'
 import { Title } from '../vscode/title'
@@ -17,7 +17,7 @@ const mockedGetFileExt = jest.mocked(getFileExtension)
 const mockedGetInput = jest.mocked(getInput)
 const mockedQuickPickOne = jest.mocked(quickPickOne)
 const mockedQuickPickValue = jest.mocked(quickPickValue)
-const mockedQuickPickValues = jest.mocked(quickPickManyValues)
+const mockedQuickPickUserOrderedValues = jest.mocked(quickPickUserOrderedValues)
 const mockedToast = jest.mocked(Toast)
 const mockedShowError = jest.fn()
 mockedToast.showError = mockedShowError
@@ -264,13 +264,14 @@ describe('pickPlotConfiguration', () => {
       { data: mockValidData, file: '/file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
-    mockedQuickPickValue.mockResolvedValueOnce({
-      file: 'file.json',
-      key: 'actual'
-    })
-    mockedQuickPickValues.mockResolvedValueOnce([
-      { file: 'file.json', key: 'prob' }
-    ])
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        }
+      ])
+      .mockResolvedValueOnce([{ file: 'file.json', key: 'prob' }])
     mockedGetInput.mockResolvedValueOnce('Simple Plot')
 
     const result = await pickPlotConfiguration('/')
@@ -290,7 +291,8 @@ describe('pickPlotConfiguration', () => {
       ],
       'Pick a Plot Template'
     )
-    expect(mockedQuickPickValue).toHaveBeenCalledWith(
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      1,
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -302,7 +304,8 @@ describe('pickPlotConfiguration', () => {
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
-    expect(mockedQuickPickValues).toHaveBeenCalledWith(
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      2,
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -338,17 +341,14 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickValue.mockResolvedValueOnce({
-      file: 'file.json',
-      key: 'actual'
-    })
-    mockedQuickPickValues.mockResolvedValueOnce([
-      { file: 'file2.json', key: 'prob' }
-    ])
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([{ file: 'file.json', key: 'actual' }])
+      .mockResolvedValueOnce([{ file: 'file2.json', key: 'prob' }])
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValue).toHaveBeenCalledWith(
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      1,
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -367,7 +367,8 @@ describe('pickPlotConfiguration', () => {
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
-    expect(mockedQuickPickValues).toHaveBeenCalledWith(
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      2,
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -406,19 +407,23 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickValue.mockResolvedValueOnce({
-      file: 'file.json',
-      key: 'actual'
-    })
-    mockedQuickPickValues.mockResolvedValueOnce([
-      { file: 'file2.json', key: 'prob' },
-      { file: 'file2.json', key: 'actual' },
-      { file: 'file2.json', key: 'step' }
-    ])
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        }
+      ])
+      .mockResolvedValueOnce([
+        { file: 'file2.json', key: 'prob' },
+        { file: 'file2.json', key: 'actual' },
+        { file: 'file2.json', key: 'step' }
+      ])
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValue).toHaveBeenCalledWith(
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      1,
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -438,7 +443,8 @@ describe('pickPlotConfiguration', () => {
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
-    expect(mockedQuickPickValues).toHaveBeenCalledWith(
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      2,
       [
         {
           kind: QuickPickItemKind.Separator,
@@ -467,6 +473,83 @@ describe('pickPlotConfiguration', () => {
     })
   })
 
+  it('should let the user pick multiple x fields', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json', '/file2.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: '/file.json' },
+      {
+        data: mockValidData,
+        file: '/file2.json'
+      }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        { file: 'file.json', key: 'prob' },
+        { file: 'file2.json', key: 'actual' }
+      ])
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        },
+        {
+          file: 'file2.json',
+          key: 'prob'
+        }
+      ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      1,
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file2.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+      ],
+      { title: Title.SELECT_PLOT_X_METRIC }
+    )
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenNthCalledWith(
+      2,
+      [
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file.json',
+          value: undefined
+        },
+        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        {
+          kind: QuickPickItemKind.Separator,
+          label: 'file2.json',
+          value: undefined
+        },
+        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+      ],
+      {
+        title: Title.SELECT_PLOT_Y_METRIC
+      }
+    )
+    expect(result).toStrictEqual({
+      template: 'simple',
+      title: 'simple_plot',
+      x: { 'file.json': 'prob', 'file2.json': 'actual' },
+      y: { 'file.json': 'actual', 'file2.json': 'prob' }
+    })
+  })
+
   it('should return early if the user does not pick a template', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
@@ -488,12 +571,37 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickValue.mockResolvedValueOnce(undefined)
+    mockedQuickPickUserOrderedValues.mockResolvedValueOnce(undefined)
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
+    expect(result).toStrictEqual(undefined)
+  })
+
+  it('should show a toast if user selects more than one key in a file for an x field', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickUserOrderedValues.mockResolvedValueOnce([
+      {
+        file: 'file.json',
+        key: 'actual'
+      },
+      {
+        file: 'file.json',
+        key: 'prob'
+      }
+    ])
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValue).toHaveBeenCalledTimes(1)
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
   })
 
   it('should return early if the user does not pick a y field', async () => {
@@ -504,30 +612,106 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickValue.mockResolvedValueOnce('actual')
-    mockedQuickPickValues.mockResolvedValueOnce(undefined)
+    mockedQuickPickUserOrderedValues.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
 
-    expect(mockedQuickPickValues).toHaveBeenCalledTimes(1)
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
   })
 
+  it('should show a toast if user selects multiple x fields and more than one key in a file for an y field', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' },
+      {
+        data: mockValidData.map((value, ind) => ({ ...value, step: ind })),
+        file: 'file2.json'
+      }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        },
+        {
+          file: 'file2.json',
+          key: 'prob'
+        }
+      ])
+      .mockResolvedValueOnce([
+        {
+          file: 'file2.json',
+          key: 'actual'
+        },
+        {
+          file: 'file2.json',
+          key: 'step'
+        }
+      ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+  })
+
+  it('should show a toast if user does not select the same amount of x and y fields (if there are multiple x fields)', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' },
+      { data: mockValidData, file: 'file2.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        },
+        {
+          file: 'file.json',
+          key: 'prob'
+        }
+      ])
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'prob'
+        }
+      ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+  })
+
   it('should return early if the user does not pick a title', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: 'file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('linear')
-    mockedQuickPickValue.mockResolvedValueOnce({
-      file: 'file.json',
-      key: 'actual'
-    })
-    mockedQuickPickValues.mockResolvedValueOnce([
-      {
-        file: 'file.json',
-        key: 'prob'
-      }
-    ])
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        }
+      ])
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'prob'
+        }
+      ])
     mockedGetInput.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 88178ea5a7..3f140090c1 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -8,11 +8,7 @@ import {
   isValueTree
 } from '../cli/dvc/contract'
 import { getFileExtension, loadDataFiles } from '../fileSystem'
-import {
-  quickPickManyValues,
-  quickPickOne,
-  quickPickValue
-} from '../vscode/quickPick'
+import { quickPickOne, quickPickUserOrderedValues } from '../vscode/quickPick'
 import { pickFiles } from '../vscode/resourcePicker'
 import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
@@ -78,16 +74,18 @@ const pickFields = async (
     )
   }
 
-  const xValue = await quickPickValue(items, {
+  const xValues = (await quickPickUserOrderedValues(items, {
     title: Title.SELECT_PLOT_X_METRIC
-  })
+  })) as { file: string; key: string }[] | undefined
 
-  if (!xValue) {
+  if (!xValues) {
     return
   }
 
-  const yValues = (await quickPickManyValues(
-    items.filter(item => !isEqual(item.value, xValue)),
+  // validate that x values are all from different files
+
+  const yValues = (await quickPickUserOrderedValues(
+    items.filter(item => xValues.every(val => !isEqual(val, item.value))),
     { title: Title.SELECT_PLOT_Y_METRIC }
   )) as { file: string; key: string }[] | undefined
 
@@ -95,12 +93,14 @@ const pickFields = async (
     return
   }
 
+  // validate that IF X VALUES IS MORE THAN 1: y values are all from different files and the same length as x values
+
   return {
     fields: {
-      x: { [xValue.file]: xValue.key },
+      x: formatFieldQuickPickValues(xValues),
       y: formatFieldQuickPickValues(yValues)
     },
-    firstXKey: xValue.key,
+    firstXKey: xValues[0].key,
     firstYKey: yValues[0].key
   }
 }
diff --git a/extension/src/vscode/quickPick.ts b/extension/src/vscode/quickPick.ts
index 430d5fd739..9f25eb5bb7 100644
--- a/extension/src/vscode/quickPick.ts
+++ b/extension/src/vscode/quickPick.ts
@@ -1,4 +1,5 @@
 import { QuickPickOptions, QuickPickItem, window, QuickPick } from 'vscode'
+import isEqual from 'lodash.isequal'
 import { Response } from './response'
 import { Title } from './title'
 
@@ -219,6 +220,45 @@ export const quickPickLimitedValues = <T>(
     quickPick.show()
   })
 
+export const quickPickUserOrderedValues = <T>(
+  items: QuickPickItemWithValue<T>[],
+  options: QuickPickOptions & { title: Title }
+): Promise<T[] | undefined> =>
+  new Promise(resolve => {
+    const quickPick = createQuickPick(items, [], {
+      ...DEFAULT_OPTIONS,
+      canSelectMany: true,
+      ...options
+    })
+
+    let orderedSelection: T[] = []
+
+    quickPick.onDidChangeSelection(selectedItems => {
+      const itemValues: T[] = []
+      for (const item of selectedItems) {
+        if (!orderedSelection.includes(item.value)) {
+          orderedSelection.push(item.value)
+        }
+        itemValues.push(item.value)
+      }
+      orderedSelection = orderedSelection.filter(item =>
+        itemValues.some(val => isEqual(item, val))
+      )
+    })
+
+    quickPick.onDidAccept(() => {
+      resolve(orderedSelection)
+      quickPick.dispose()
+    })
+
+    quickPick.onDidHide(() => {
+      resolve(undefined)
+      quickPick.dispose()
+    })
+
+    quickPick.show()
+  })
+
 export const quickPickYesOrNo = (
   descriptionYes: string,
   descriptionNo: string,

From 5933faac153a09167478dc3fa329f8906864cee8 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 10 Oct 2023 11:52:25 -0500
Subject: [PATCH 27/35] Move error handling to chain pr

---
 extension/src/pipeline/quickPick.test.ts | 99 ------------------------
 extension/src/pipeline/quickPick.ts      |  6 +-
 2 files changed, 1 insertion(+), 104 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 1daeeeeade..110742e76a 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -578,32 +578,6 @@ describe('pickPlotConfiguration', () => {
     expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
     expect(result).toStrictEqual(undefined)
   })
-
-  it('should show a toast if user selects more than one key in a file for an x field', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
-    mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: mockValidData, file: 'file.json' }
-    ])
-    mockedQuickPickOne.mockResolvedValueOnce('simple')
-    mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickUserOrderedValues.mockResolvedValueOnce([
-      {
-        file: 'file.json',
-        key: 'actual'
-      },
-      {
-        file: 'file.json',
-        key: 'prob'
-      }
-    ])
-
-    const result = await pickPlotConfiguration('/')
-
-    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
-    expect(result).toStrictEqual(undefined)
-    expect(mockedShowError).toHaveBeenCalledTimes(1)
-  })
-
   it('should return early if the user does not pick a y field', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
@@ -620,79 +594,6 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual(undefined)
   })
 
-  it('should show a toast if user selects multiple x fields and more than one key in a file for an y field', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
-    mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: mockValidData, file: 'file.json' },
-      {
-        data: mockValidData.map((value, ind) => ({ ...value, step: ind })),
-        file: 'file2.json'
-      }
-    ])
-    mockedQuickPickOne.mockResolvedValueOnce('simple')
-    mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickUserOrderedValues
-      .mockResolvedValueOnce([
-        {
-          file: 'file.json',
-          key: 'actual'
-        },
-        {
-          file: 'file2.json',
-          key: 'prob'
-        }
-      ])
-      .mockResolvedValueOnce([
-        {
-          file: 'file2.json',
-          key: 'actual'
-        },
-        {
-          file: 'file2.json',
-          key: 'step'
-        }
-      ])
-
-    const result = await pickPlotConfiguration('/')
-
-    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
-    expect(result).toStrictEqual(undefined)
-    expect(mockedShowError).toHaveBeenCalledTimes(1)
-  })
-
-  it('should show a toast if user does not select the same amount of x and y fields (if there are multiple x fields)', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
-    mockedLoadDataFiles.mockResolvedValueOnce([
-      { data: mockValidData, file: 'file.json' },
-      { data: mockValidData, file: 'file2.json' }
-    ])
-    mockedQuickPickOne.mockResolvedValueOnce('simple')
-    mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickUserOrderedValues
-      .mockResolvedValueOnce([
-        {
-          file: 'file.json',
-          key: 'actual'
-        },
-        {
-          file: 'file.json',
-          key: 'prob'
-        }
-      ])
-      .mockResolvedValueOnce([
-        {
-          file: 'file.json',
-          key: 'prob'
-        }
-      ])
-
-    const result = await pickPlotConfiguration('/')
-
-    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
-    expect(result).toStrictEqual(undefined)
-    expect(mockedShowError).toHaveBeenCalledTimes(1)
-  })
-
   it('should return early if the user does not pick a title', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 3f140090c1..29c6dac3f8 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -82,8 +82,6 @@ const pickFields = async (
     return
   }
 
-  // validate that x values are all from different files
-
   const yValues = (await quickPickUserOrderedValues(
     items.filter(item => xValues.every(val => !isEqual(val, item.value))),
     { title: Title.SELECT_PLOT_Y_METRIC }
@@ -93,11 +91,9 @@ const pickFields = async (
     return
   }
 
-  // validate that IF X VALUES IS MORE THAN 1: y values are all from different files and the same length as x values
-
   return {
     fields: {
-      x: formatFieldQuickPickValues(xValues),
+      x: formatFieldQuickPickValues(xValues) as PlotConfigData['x'],
       y: formatFieldQuickPickValues(yValues)
     },
     firstXKey: xValues[0].key,

From 135582a1b0ff3d84b8f36165889a4f88774fb0aa Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Tue, 10 Oct 2023 13:51:29 -0500
Subject: [PATCH 28/35] update quick pick util

---
 extension/src/vscode/quickPick.ts | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/extension/src/vscode/quickPick.ts b/extension/src/vscode/quickPick.ts
index 9f25eb5bb7..b198c1a36e 100644
--- a/extension/src/vscode/quickPick.ts
+++ b/extension/src/vscode/quickPick.ts
@@ -220,6 +220,22 @@ export const quickPickLimitedValues = <T>(
     quickPick.show()
   })
 
+const getUserOrderedSelection = <T>(
+  orderedSelection: T[],
+  newSelection: readonly QuickPickItemWithValue<T>[]
+) => {
+  const itemValues: T[] = []
+  for (const item of newSelection) {
+    if (!orderedSelection.includes(item.value)) {
+      orderedSelection.push(item.value)
+    }
+    itemValues.push(item.value)
+  }
+  return orderedSelection.filter(item =>
+    itemValues.some(val => isEqual(item, val))
+  )
+}
+
 export const quickPickUserOrderedValues = <T>(
   items: QuickPickItemWithValue<T>[],
   options: QuickPickOptions & { title: Title }
@@ -234,20 +250,14 @@ export const quickPickUserOrderedValues = <T>(
     let orderedSelection: T[] = []
 
     quickPick.onDidChangeSelection(selectedItems => {
-      const itemValues: T[] = []
-      for (const item of selectedItems) {
-        if (!orderedSelection.includes(item.value)) {
-          orderedSelection.push(item.value)
-        }
-        itemValues.push(item.value)
-      }
-      orderedSelection = orderedSelection.filter(item =>
-        itemValues.some(val => isEqual(item, val))
+      orderedSelection = getUserOrderedSelection(
+        orderedSelection,
+        selectedItems
       )
     })
 
     quickPick.onDidAccept(() => {
-      resolve(orderedSelection)
+      resolve(orderedSelection.length === 0 ? undefined : orderedSelection)
       quickPick.dispose()
     })
 

From 973e684a38d9cf769d1ac3e921a3498a035da2f3 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 11 Oct 2023 08:37:28 -0500
Subject: [PATCH 29/35] Improve types

---
 extension/src/pipeline/quickPick.ts | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index a2eaa9e6ec..9bdd80d356 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -15,7 +15,7 @@ import { Toast } from '../vscode/toast'
 import { getInput } from '../vscode/inputBox'
 
 export type PlotConfigData = {
-  x: { [file: string]: string }
+  x: { [file: string]: string[] | string }
   template: string
   title: string
   y: { [file: string]: string[] | string }
@@ -33,7 +33,7 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
 const formatFieldQuickPickValues = (
   values: { file: string; key: string }[]
 ) => {
-  const formattedFields: PlotConfigData['y'] = {}
+  const formattedFields: { [file: string]: string[] | string } = {}
 
   for (const { file, key } of values) {
     if (!formattedFields[file]) {
@@ -94,7 +94,7 @@ const pickFields = async (
 
   return {
     fields: {
-      x: formatFieldQuickPickValues(xValues) as PlotConfigData['x'],
+      x: formatFieldQuickPickValues(xValues),
       y: formatFieldQuickPickValues(yValues)
     },
     firstXKey: xValues[0].key,

From 749c2f926148685137523471e7397bb1e2cff4bd Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 11 Oct 2023 11:20:15 -0500
Subject: [PATCH 30/35] Add tests for new quick pick util

---
 extension/src/test/suite/util.ts              | 16 ++++++
 .../src/test/suite/vscode/quickPick.test.ts   | 51 ++++++++++++++++++-
 extension/src/vscode/quickPick.ts             |  2 +-
 3 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/extension/src/test/suite/util.ts b/extension/src/test/suite/util.ts
index bc7fe1ec24..a2fa93621e 100644
--- a/extension/src/test/suite/util.ts
+++ b/extension/src/test/suite/util.ts
@@ -85,6 +85,22 @@ export const selectQuickPickItem = async (number: number) => {
   return commands.executeCommand('workbench.action.acceptSelectedQuickOpenItem')
 }
 
+export const selectQuickPickItems = async (
+  numbers: number[],
+  itemsLength: number
+) => {
+  for (const number of numbers) {
+    for (let itemInd = 1; itemInd <= itemsLength; itemInd++) {
+      await commands.executeCommand('workbench.action.quickOpenSelectNext')
+
+      if (itemInd === number) {
+        await commands.executeCommand('workbench.action.quickPickManyToggle')
+      }
+    }
+  }
+  return commands.executeCommand('workbench.action.acceptSelectedQuickOpenItem')
+}
+
 export const experimentsUpdatedEvent = (experiments: Experiments) =>
   new Promise(resolve => {
     experiments.dispose.track(experiments.onDidChangeExperiments(resolve))
diff --git a/extension/src/test/suite/vscode/quickPick.test.ts b/extension/src/test/suite/vscode/quickPick.test.ts
index 08d43541dd..7e344213ed 100644
--- a/extension/src/test/suite/vscode/quickPick.test.ts
+++ b/extension/src/test/suite/vscode/quickPick.test.ts
@@ -6,9 +6,10 @@ import { Disposable } from '../../../extension'
 import {
   QuickPickItemWithValue,
   quickPickLimitedValues,
-  quickPickOneOrInput
+  quickPickOneOrInput,
+  quickPickUserOrderedValues
 } from '../../../vscode/quickPick'
-import { selectQuickPickItem } from '../util'
+import { selectQuickPickItem, selectQuickPickItems } from '../util'
 import { Title } from '../../../vscode/title'
 
 suite('Quick Pick Test Suite', () => {
@@ -140,4 +141,50 @@ suite('Quick Pick Test Suite', () => {
       expect(quickPick.items, 'all items are shown').to.deep.equal(items)
     })
   })
+
+  describe('quickPickUserOrderedValues', () => {
+    it('should return selected values in the order that the user selected', async () => {
+      const items = [
+        { label: 'J', value: 1 },
+        { label: 'K', value: 2 },
+        { label: 'L', value: 3 },
+        { label: 'M', value: 4 },
+        { label: 'N', value: 5 }
+      ]
+
+      const resultPromise = quickPickUserOrderedValues(items, {
+        title: 'Select some values' as Title
+      })
+
+      await selectQuickPickItems([5, 2, 1], items.length)
+
+      const result = await resultPromise
+
+      expect(result).to.deep.equal([5, 2, 1])
+    })
+
+    it('should return undefined if user cancels the quick pick', async () => {
+      const quickPick = disposable.track(
+        window.createQuickPick<QuickPickItemWithValue<number>>()
+      )
+      stub(window, 'createQuickPick').returns(quickPick)
+      const items = [
+        { label: 'J', value: 1 },
+        { label: 'K', value: 2 },
+        { label: 'L', value: 3 },
+        { label: 'M', value: 4 },
+        { label: 'N', value: 5 }
+      ]
+
+      const resultPromise = quickPickUserOrderedValues(items, {
+        title: 'Select some values' as Title
+      })
+
+      quickPick.hide()
+
+      const result = await resultPromise
+
+      expect(result).to.deep.equal(undefined)
+    })
+  })
 })
diff --git a/extension/src/vscode/quickPick.ts b/extension/src/vscode/quickPick.ts
index b198c1a36e..a394c6ba50 100644
--- a/extension/src/vscode/quickPick.ts
+++ b/extension/src/vscode/quickPick.ts
@@ -257,7 +257,7 @@ export const quickPickUserOrderedValues = <T>(
     })
 
     quickPick.onDidAccept(() => {
-      resolve(orderedSelection.length === 0 ? undefined : orderedSelection)
+      resolve(orderedSelection)
       quickPick.dispose()
     })
 

From 880a6c04d3992fdd1c240f36479e6c76f40baaee Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 11 Oct 2023 11:25:55 -0500
Subject: [PATCH 31/35] Add check for empty x fields

---
 extension/src/pipeline/quickPick.test.ts | 21 ++++++++++++++-------
 extension/src/pipeline/quickPick.ts      |  2 +-
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 110742e76a..1732260657 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -565,18 +565,25 @@ describe('pickPlotConfiguration', () => {
   })
 
   it('should return early if the user does not pick a x field', async () => {
-    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
-    mockedLoadDataFiles.mockResolvedValueOnce([
+    mockedPickFiles.mockResolvedValue(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValue([
       { data: mockValidData, file: 'file.json' }
     ])
-    mockedQuickPickOne.mockResolvedValueOnce('simple')
-    mockedGetInput.mockResolvedValueOnce('simple_plot')
-    mockedQuickPickUserOrderedValues.mockResolvedValueOnce(undefined)
+    mockedQuickPickOne.mockResolvedValue('simple')
+    mockedGetInput.mockResolvedValue('simple_plot')
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce(undefined)
+      .mockResolvedValueOnce([])
 
-    const result = await pickPlotConfiguration('/')
+    const undefinedResult = await pickPlotConfiguration('/')
 
     expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
-    expect(result).toStrictEqual(undefined)
+    expect(undefinedResult).toStrictEqual(undefined)
+
+    const noFieldsResult = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
+    expect(noFieldsResult).toStrictEqual(undefined)
   })
   it('should return early if the user does not pick a y field', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index 9bdd80d356..5b27d4f21f 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -79,7 +79,7 @@ const pickFields = async (
     title: Title.SELECT_PLOT_X_METRIC
   })) as { file: string; key: string }[] | undefined
 
-  if (!xValues) {
+  if (!xValues || xValues.length === 0) {
     return
   }
 

From c6349bb8c8f76ad7f40d893ac2aa8b928de7138b Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Wed, 11 Oct 2023 11:45:50 -0500
Subject: [PATCH 32/35] Minor refactor

---
 extension/src/fileSystem/index.ts | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index 11f16862e6..2d53adc61b 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -219,13 +219,14 @@ const getPlotYamlObj = (plot: PlotConfigData) => {
 
   const yFiles = Object.keys(y)
   const xFiles = Object.keys(x)
+  const firstXFile = xFiles[0]
   const oneFileUsed =
-    yFiles.length === 1 && xFiles.length === 1 && yFiles[0] === xFiles[0]
+    yFiles.length === 1 && xFiles.length === 1 && yFiles[0] === firstXFile
 
   return {
     [title]: {
       template,
-      x: oneFileUsed ? x[xFiles[0]] : x,
+      x: oneFileUsed ? x[firstXFile] : x,
       y
     }
   }

From 1d892f8f158f5e776f42d59d52eb218d4c3205e3 Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 12 Oct 2023 12:05:36 -0500
Subject: [PATCH 33/35] Resolve review comments

---
 extension/src/fileSystem/index.test.ts        |  50 +++++--
 extension/src/fileSystem/index.ts             |  21 ++-
 extension/src/pipeline/quickPick.test.ts      | 122 ++++++++----------
 extension/src/pipeline/quickPick.ts           |  45 +++----
 .../src/test/suite/pipeline/index.test.ts     |   4 +-
 extension/src/test/suite/util.ts              |  20 +--
 .../src/test/suite/vscode/quickPick.test.ts   |   4 +-
 extension/src/test/util/mocha/index.ts        |   2 +-
 extension/src/vscode/quickPick.ts             |  20 ++-
 9 files changed, 163 insertions(+), 125 deletions(-)

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
index 0cd03c57ed..02a2715f21 100644
--- a/extension/src/fileSystem/index.test.ts
+++ b/extension/src/fileSystem/index.test.ts
@@ -625,8 +625,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -654,8 +654,38 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'simple_plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'acc.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'acc.json': ['accuracy'] }
+    })
+
+    expect(mockedWriteFileSync).toHaveBeenCalledWith(
+      '//dvc.yaml',
+      mockDvcYamlContent + mockPlotYamlContent
+    )
+  })
+
+  it('should add the new plot with an axis having multiple fields', () => {
+    const mockDvcYamlContent = mockStagesLines.join('\n')
+    const mockPlotYamlContent = [
+      '',
+      'plots:',
+      '  - simple_plot:',
+      '      template: simple',
+      '      x: epochs',
+      '      y:',
+      '        data.json:',
+      '          - accuracy',
+      '          - epochs',
+      ''
+    ].join('\n')
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+    mockedReadFileSync.mockReturnValueOnce(mockDvcYamlContent)
+
+    addPlotToDvcYamlFile('/', {
+      template: 'simple',
+      title: 'simple_plot',
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy', 'epochs'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -673,8 +703,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     mockDvcYamlContent.splice(7, 0, ...mockPlotYamlContent)
@@ -696,8 +726,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'Simple Plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
@@ -734,8 +764,8 @@ describe('addPlotToDvcYamlFile', () => {
     addPlotToDvcYamlFile('/', {
       template: 'simple',
       title: 'simple_plot',
-      x: { 'data.json': 'epochs' },
-      y: { 'data.json': 'accuracy' }
+      x: { 'data.json': ['epochs'] },
+      y: { 'data.json': ['accuracy'] }
     })
 
     expect(mockedWriteFileSync).toHaveBeenCalledWith(
diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
index 2d53adc61b..59965c13a4 100644
--- a/extension/src/fileSystem/index.ts
+++ b/extension/src/fileSystem/index.ts
@@ -30,7 +30,7 @@ import { processExists } from '../process/execution'
 import { getFirstWorkspaceFolder } from '../vscode/workspaceFolders'
 import { DOT_DVC } from '../cli/dvc/constants'
 import { delay } from '../util/time'
-import { PlotConfigData } from '../pipeline/quickPick'
+import { PlotConfigData, PlotConfigDataAxis } from '../pipeline/quickPick'
 
 export const exists = (path: string): boolean => existsSync(path)
 
@@ -214,6 +214,21 @@ const loadYamlAsDoc = (
   }
 }
 
+const formatPlotYamlObjAxis = (axis: PlotConfigDataAxis) => {
+  const formattedAxis: { [file: string]: string | string[] } = {}
+
+  for (const [file, fields] of Object.entries(axis)) {
+    if (fields.length === 1) {
+      formattedAxis[file] = fields[0]
+      continue
+    }
+
+    formattedAxis[file] = fields
+  }
+
+  return formattedAxis
+}
+
 const getPlotYamlObj = (plot: PlotConfigData) => {
   const { x, y, template, title } = plot
 
@@ -226,8 +241,8 @@ const getPlotYamlObj = (plot: PlotConfigData) => {
   return {
     [title]: {
       template,
-      x: oneFileUsed ? x[firstXFile] : x,
-      y
+      x: oneFileUsed ? x[firstXFile][0] : formatPlotYamlObjAxis(x),
+      y: formatPlotYamlObjAxis(y)
     }
   }
 }
diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index 30d4db6735..bc60922f56 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -260,13 +260,8 @@ describe('pickPlotConfiguration', () => {
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedQuickPickUserOrderedValues
-      .mockResolvedValueOnce([
-        {
-          file: 'file.json',
-          key: 'actual'
-        }
-      ])
-      .mockResolvedValueOnce([{ file: 'file.json', key: 'prob' }])
+      .mockImplementationOnce(items => Promise.resolve([items[1].value]))
+      .mockImplementationOnce(items => Promise.resolve([items[1].value]))
     mockedGetInput.mockResolvedValueOnce('Simple Plot')
 
     const result = await pickPlotConfiguration('/')
@@ -294,8 +289,8 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
+        { label: 'actual', value: { field: 'actual', file: 'file.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file.json' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -307,7 +302,7 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } }
+        { label: 'prob', value: { field: 'prob', file: 'file.json' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -320,8 +315,8 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'Simple Plot',
-      x: { 'file.json': 'actual' },
-      y: { 'file.json': 'prob' }
+      x: { 'file.json': ['actual'] },
+      y: { 'file.json': ['prob'] }
     })
   })
 
@@ -337,8 +332,8 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickUserOrderedValues
-      .mockResolvedValueOnce([{ file: 'file.json', key: 'actual' }])
-      .mockResolvedValueOnce([{ file: 'file2.json', key: 'prob' }])
+      .mockImplementationOnce(items => Promise.resolve([items[1].value]))
+      .mockImplementationOnce(items => Promise.resolve([items[4].value]))
 
     const result = await pickPlotConfiguration('/')
 
@@ -350,15 +345,15 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        { label: 'actual', value: { field: 'actual', file: 'file.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file.json' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'actual', value: { field: 'actual', file: 'file2.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file2.json' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -370,14 +365,14 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        { label: 'prob', value: { field: 'prob', file: 'file.json' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'actual', value: { field: 'actual', file: 'file2.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file2.json' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -386,8 +381,8 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'simple_plot',
-      x: { 'file.json': 'actual' },
-      y: { 'file2.json': 'prob' }
+      x: { 'file.json': ['actual'] },
+      y: { 'file2.json': ['prob'] }
     })
   })
 
@@ -403,17 +398,10 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickUserOrderedValues
-      .mockResolvedValueOnce([
-        {
-          file: 'file.json',
-          key: 'actual'
-        }
-      ])
-      .mockResolvedValueOnce([
-        { file: 'file2.json', key: 'prob' },
-        { file: 'file2.json', key: 'actual' },
-        { file: 'file2.json', key: 'step' }
-      ])
+      .mockImplementationOnce(items => Promise.resolve([items[1].value]))
+      .mockImplementationOnce(items =>
+        Promise.resolve([items[3].value, items[4].value, items[5].value])
+      )
 
     const result = await pickPlotConfiguration('/')
 
@@ -425,16 +413,16 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        { label: 'actual', value: { field: 'actual', file: 'file.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file.json' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } },
-        { label: 'step', value: { file: 'file2.json', key: 'step' } }
+        { label: 'actual', value: { field: 'actual', file: 'file2.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file2.json' } },
+        { label: 'step', value: { field: 'step', file: 'file2.json' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -446,15 +434,15 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        { label: 'prob', value: { field: 'prob', file: 'file.json' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } },
-        { label: 'step', value: { file: 'file2.json', key: 'step' } }
+        { label: 'actual', value: { field: 'actual', file: 'file2.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file2.json' } },
+        { label: 'step', value: { field: 'step', file: 'file2.json' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -463,8 +451,8 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'simple_plot',
-      x: { 'file.json': 'actual' },
-      y: { 'file2.json': ['prob', 'actual', 'step'] }
+      x: { 'file.json': ['actual'] },
+      y: { 'file2.json': ['actual', 'prob', 'step'] }
     })
   })
 
@@ -480,20 +468,12 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickOne.mockResolvedValueOnce('simple')
     mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickUserOrderedValues
-      .mockResolvedValueOnce([
-        { file: 'file.json', key: 'prob' },
-        { file: 'file2.json', key: 'actual' }
-      ])
-      .mockResolvedValueOnce([
-        {
-          file: 'file.json',
-          key: 'actual'
-        },
-        {
-          file: 'file2.json',
-          key: 'prob'
-        }
-      ])
+      .mockImplementationOnce(items =>
+        Promise.resolve([items[2].value, items[4].value])
+      )
+      .mockImplementationOnce(items =>
+        Promise.resolve([items[1].value, items[3].value])
+      )
 
     const result = await pickPlotConfiguration('/')
 
@@ -505,15 +485,15 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file.json', key: 'prob' } },
+        { label: 'actual', value: { field: 'actual', file: 'file.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file.json' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file2.json', key: 'actual' } },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'actual', value: { field: 'actual', file: 'file2.json' } },
+        { label: 'prob', value: { field: 'prob', file: 'file2.json' } }
       ],
       { title: Title.SELECT_PLOT_X_METRIC }
     )
@@ -525,13 +505,13 @@ describe('pickPlotConfiguration', () => {
           label: 'file.json',
           value: undefined
         },
-        { label: 'actual', value: { file: 'file.json', key: 'actual' } },
+        { label: 'actual', value: { field: 'actual', file: 'file.json' } },
         {
           kind: QuickPickItemKind.Separator,
           label: 'file2.json',
           value: undefined
         },
-        { label: 'prob', value: { file: 'file2.json', key: 'prob' } }
+        { label: 'prob', value: { field: 'prob', file: 'file2.json' } }
       ],
       {
         title: Title.SELECT_PLOT_Y_METRIC
@@ -540,8 +520,8 @@ describe('pickPlotConfiguration', () => {
     expect(result).toStrictEqual({
       template: 'simple',
       title: 'simple_plot',
-      x: { 'file.json': 'prob', 'file2.json': 'actual' },
-      y: { 'file.json': 'actual', 'file2.json': 'prob' }
+      x: { 'file.json': ['prob'], 'file2.json': ['actual'] },
+      y: { 'file.json': ['actual'], 'file2.json': ['prob'] }
     })
   })
 
@@ -610,14 +590,14 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickUserOrderedValues
       .mockResolvedValueOnce([
         {
-          file: 'file.json',
-          key: 'actual'
+          field: 'actual',
+          file: 'file.json'
         }
       ])
       .mockResolvedValueOnce([
         {
-          file: 'file.json',
-          key: 'prob'
+          field: 'prob',
+          file: 'file.json'
         }
       ])
     mockedGetInput.mockResolvedValueOnce(undefined)
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index c3fe12eba4..b84125b3ee 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -14,11 +14,13 @@ import { Title } from '../vscode/title'
 import { Toast } from '../vscode/toast'
 import { getInput } from '../vscode/inputBox'
 
+export type PlotConfigDataAxis = { [file: string]: string[] }
+
 export type PlotConfigData = {
-  x: { [file: string]: string[] | string }
+  x: PlotConfigDataAxis
   template: string
   title: string
-  y: { [file: string]: string[] | string }
+  y: PlotConfigDataAxis
 }
 
 type UnknownValue = Value | ValueTree
@@ -31,22 +33,17 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
   })
 
 const formatFieldQuickPickValues = (
-  values: { file: string; key: string }[]
+  values: { file: string; field: string }[]
 ) => {
-  const formattedFields: { [file: string]: string[] | string } = {}
+  const formattedFields: PlotConfigDataAxis = {}
 
-  for (const { file, key } of values) {
+  for (const { file, field } of values) {
     if (!formattedFields[file]) {
-      formattedFields[file] = key
+      formattedFields[file] = [field]
       continue
     }
 
-    const prevFileValue = formattedFields[file]
-    if (typeof prevFileValue === 'string') {
-      formattedFields[file] = [prevFileValue]
-    }
-
-    ;(formattedFields[file] as string[]).push(key)
+    formattedFields[file].push(field)
   }
 
   return formattedFields
@@ -71,22 +68,26 @@ const pickFields = async (
         label: file,
         value: undefined
       },
-      ...fields.map(key => ({ label: key, value: { file, key } }))
+      ...fields.map(field => ({ label: field, value: { field, file } }))
     )
   }
 
   const xValues = (await quickPickUserOrderedValues(items, {
     title: Title.SELECT_PLOT_X_METRIC
-  })) as { file: string; key: string }[] | undefined
+  })) as { file: string; field: string }[] | undefined
 
   if (!xValues || xValues.length === 0) {
     return
   }
 
   const yValues = (await quickPickUserOrderedValues(
-    items.filter(item => xValues.every(val => !isEqual(val, item.value))),
-    { title: Title.SELECT_PLOT_Y_METRIC }
-  )) as { file: string; key: string }[] | undefined
+    items.filter(
+      item => item.value === undefined || !xValues.includes(item.value)
+    ),
+    {
+      title: Title.SELECT_PLOT_Y_METRIC
+    }
+  )) as { file: string; field: string }[] | undefined
 
   if (!yValues || yValues.length === 0) {
     return
@@ -97,8 +98,8 @@ const pickFields = async (
       x: formatFieldQuickPickValues(xValues),
       y: formatFieldQuickPickValues(yValues)
     },
-    firstXKey: xValues[0].key,
-    firstYKey: yValues[0].key
+    firstXKey: xValues[0].field,
+    firstYKey: yValues[0].field
   }
 }
 
@@ -220,7 +221,7 @@ const validateMultiFilesData = (
   filesData: { data: UnknownValue; file: string }[]
 ) => {
   const filesArrLength: Set<number> = new Set()
-  const keys: FileFields = []
+  const fileFields: FileFields = []
 
   for (const { file, data } of filesData) {
     const metricInfo = getMetricInfoFromValue(data)
@@ -230,7 +231,7 @@ const validateMultiFilesData = (
     }
 
     const { arrLength, fields } = metricInfo
-    keys.push({ fields, file })
+    fileFields.push({ fields, file })
     filesArrLength.add(arrLength)
   }
 
@@ -242,7 +243,7 @@ const validateMultiFilesData = (
     return
   }
 
-  return keys
+  return fileFields
 }
 
 const validateFilesData = async (cwd: string, files: string[]) => {
diff --git a/extension/src/test/suite/pipeline/index.test.ts b/extension/src/test/suite/pipeline/index.test.ts
index 027fe97012..116164668a 100644
--- a/extension/src/test/suite/pipeline/index.test.ts
+++ b/extension/src/test/suite/pipeline/index.test.ts
@@ -312,8 +312,8 @@ suite('Pipeline Test Suite', () => {
     const mockPlotConfig: PipelineQuickPick.PlotConfigData = {
       template: 'simple',
       title: 'Great Plot Name',
-      x: { 'results.json': 'step' },
-      y: { 'results.json': 'acc' }
+      x: { 'results.json': ['step'] },
+      y: { 'results.json': ['acc'] }
     }
 
     mockPickPlotConfiguration.onSecondCall().resolves(mockPlotConfig)
diff --git a/extension/src/test/suite/util.ts b/extension/src/test/suite/util.ts
index a2fa93621e..55269f090b 100644
--- a/extension/src/test/suite/util.ts
+++ b/extension/src/test/suite/util.ts
@@ -85,18 +85,22 @@ export const selectQuickPickItem = async (number: number) => {
   return commands.executeCommand('workbench.action.acceptSelectedQuickOpenItem')
 }
 
-export const selectQuickPickItems = async (
+const toggleQuickPickItem = async (number: number, itemsLength: number) => {
+  for (let itemInd = 1; itemInd <= itemsLength; itemInd++) {
+    await commands.executeCommand('workbench.action.quickOpenSelectNext')
+
+    if (itemInd === number) {
+      await commands.executeCommand('workbench.action.quickPickManyToggle')
+    }
+  }
+}
+
+export const selectMultipleQuickPickItems = async (
   numbers: number[],
   itemsLength: number
 ) => {
   for (const number of numbers) {
-    for (let itemInd = 1; itemInd <= itemsLength; itemInd++) {
-      await commands.executeCommand('workbench.action.quickOpenSelectNext')
-
-      if (itemInd === number) {
-        await commands.executeCommand('workbench.action.quickPickManyToggle')
-      }
-    }
+    await toggleQuickPickItem(number, itemsLength)
   }
   return commands.executeCommand('workbench.action.acceptSelectedQuickOpenItem')
 }
diff --git a/extension/src/test/suite/vscode/quickPick.test.ts b/extension/src/test/suite/vscode/quickPick.test.ts
index 7e344213ed..85b58e61e9 100644
--- a/extension/src/test/suite/vscode/quickPick.test.ts
+++ b/extension/src/test/suite/vscode/quickPick.test.ts
@@ -9,7 +9,7 @@ import {
   quickPickOneOrInput,
   quickPickUserOrderedValues
 } from '../../../vscode/quickPick'
-import { selectQuickPickItem, selectQuickPickItems } from '../util'
+import { selectQuickPickItem, selectMultipleQuickPickItems } from '../util'
 import { Title } from '../../../vscode/title'
 
 suite('Quick Pick Test Suite', () => {
@@ -156,7 +156,7 @@ suite('Quick Pick Test Suite', () => {
         title: 'Select some values' as Title
       })
 
-      await selectQuickPickItems([5, 2, 1], items.length)
+      await selectMultipleQuickPickItems([5, 2, 1], items.length)
 
       const result = await resultPromise
 
diff --git a/extension/src/test/util/mocha/index.ts b/extension/src/test/util/mocha/index.ts
index 1b49a01bb4..680f91a9e8 100644
--- a/extension/src/test/util/mocha/index.ts
+++ b/extension/src/test/util/mocha/index.ts
@@ -29,7 +29,7 @@ export const runMocha = async (
   try {
     // eslint-disable-next-line sonarjs/cognitive-complexity
     await new Promise<void>((resolve, reject) => {
-      glob(`**/**.test.${ext}`, { cwd }, (err, files) => {
+      glob(`**/quickPick.test.${ext}`, { cwd }, (err, files) => {
         if (err) {
           return reject(err)
         }
diff --git a/extension/src/vscode/quickPick.ts b/extension/src/vscode/quickPick.ts
index a394c6ba50..1b36496675 100644
--- a/extension/src/vscode/quickPick.ts
+++ b/extension/src/vscode/quickPick.ts
@@ -1,5 +1,4 @@
 import { QuickPickOptions, QuickPickItem, window, QuickPick } from 'vscode'
-import isEqual from 'lodash.isequal'
 import { Response } from './response'
 import { Title } from './title'
 
@@ -220,7 +219,7 @@ export const quickPickLimitedValues = <T>(
     quickPick.show()
   })
 
-const getUserOrderedSelection = <T>(
+const addNewUserOrderedItems = <T>(
   orderedSelection: T[],
   newSelection: readonly QuickPickItemWithValue<T>[]
 ) => {
@@ -231,11 +230,15 @@ const getUserOrderedSelection = <T>(
     }
     itemValues.push(item.value)
   }
-  return orderedSelection.filter(item =>
-    itemValues.some(val => isEqual(item, val))
-  )
+
+  return itemValues
 }
 
+const removeMissingUserOrderedItems = <T>(
+  orderedSelection: T[],
+  newSelection: T[]
+) => orderedSelection.filter(item => newSelection.includes(item))
+
 export const quickPickUserOrderedValues = <T>(
   items: QuickPickItemWithValue<T>[],
   options: QuickPickOptions & { title: Title }
@@ -250,10 +253,15 @@ export const quickPickUserOrderedValues = <T>(
     let orderedSelection: T[] = []
 
     quickPick.onDidChangeSelection(selectedItems => {
-      orderedSelection = getUserOrderedSelection(
+      const selectedItemValues = addNewUserOrderedItems(
         orderedSelection,
         selectedItems
       )
+
+      orderedSelection = removeMissingUserOrderedItems(
+        orderedSelection,
+        selectedItemValues
+      )
     })
 
     quickPick.onDidAccept(() => {

From 974126081efaf070546e8454dc71a6db327998ad Mon Sep 17 00:00:00 2001
From: julieg18 <juliannagal18@gmail.com>
Date: Thu, 12 Oct 2023 12:07:46 -0500
Subject: [PATCH 34/35] Fix typo

---
 extension/src/test/util/mocha/index.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/extension/src/test/util/mocha/index.ts b/extension/src/test/util/mocha/index.ts
index 680f91a9e8..1b49a01bb4 100644
--- a/extension/src/test/util/mocha/index.ts
+++ b/extension/src/test/util/mocha/index.ts
@@ -29,7 +29,7 @@ export const runMocha = async (
   try {
     // eslint-disable-next-line sonarjs/cognitive-complexity
     await new Promise<void>((resolve, reject) => {
-      glob(`**/quickPick.test.${ext}`, { cwd }, (err, files) => {
+      glob(`**/**.test.${ext}`, { cwd }, (err, files) => {
         if (err) {
           return reject(err)
         }

From 7033ec674e540597e32e6573c885ee30a59ba78c Mon Sep 17 00:00:00 2001
From: Julie G <43496356+julieg18@users.noreply.github.com>
Date: Fri, 13 Oct 2023 08:51:52 -0500
Subject: [PATCH 35/35] Add error handling for Plot Wizard plots with x fields
 (#4798)

---
 extension/src/pipeline/quickPick.test.ts | 127 ++++++++++++++++++++++-
 extension/src/pipeline/quickPick.ts      | 101 ++++++++++++++++--
 2 files changed, 214 insertions(+), 14 deletions(-)

diff --git a/extension/src/pipeline/quickPick.test.ts b/extension/src/pipeline/quickPick.test.ts
index bc60922f56..79715e0361 100644
--- a/extension/src/pipeline/quickPick.test.ts
+++ b/extension/src/pipeline/quickPick.test.ts
@@ -560,6 +560,35 @@ describe('pickPlotConfiguration', () => {
     expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
     expect(noFieldsResult).toStrictEqual(undefined)
   })
+
+  it('should show a toast if user selects more than one key in a file for an x field', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickUserOrderedValues.mockResolvedValueOnce([
+      {
+        file: 'file.json',
+        key: 'actual'
+      },
+      {
+        file: 'file.json',
+        key: 'prob'
+      }
+    ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+    expect(mockedShowError).toHaveBeenCalledWith(
+      'file.json cannot have more than one metric selected. See [an example](https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#available-configuration-fields) of a plot with multiple x metrics.'
+    )
+  })
+
   it('should return early if the user does not pick a y field', async () => {
     mockedPickFiles.mockResolvedValue(['/file.json'])
     mockedLoadDataFiles.mockResolvedValue([
@@ -568,17 +597,109 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickOne.mockResolvedValue('simple')
     mockedGetInput.mockResolvedValue('simple_plot')
     mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        }
+      ])
       .mockResolvedValueOnce(undefined)
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        }
+      ])
       .mockResolvedValueOnce([])
 
     const undefinedResult = await pickPlotConfiguration('/')
-    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(1)
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
     expect(undefinedResult).toStrictEqual(undefined)
 
-    const emptyFieldsResult = await pickPlotConfiguration('/')
+    const noFieldsResult = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(4)
+    expect(noFieldsResult).toStrictEqual(undefined)
+  })
+
+  it('should show a toast if user does not select the same amount of x and y fields (if there are multiple x fields)', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' },
+      { data: mockValidData, file: 'file2.json' }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        },
+        {
+          file: 'file2.json',
+          key: 'prob'
+        }
+      ])
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'prob'
+        }
+      ])
+
+    const result = await pickPlotConfiguration('/')
 
     expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
-    expect(emptyFieldsResult).toStrictEqual(undefined)
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+    expect(mockedShowError).toHaveBeenCalledWith(
+      'Found 2 x metrics and 1 y metric(s). When there are multiple x metrics selected there must be an equal number of y metrics. See [an example](https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#available-configuration-fields) of a plot with multiple x metrics.'
+    )
+  })
+
+  it('should show a toast if user selects multiple x fields and more than one key in a file for an y field', async () => {
+    mockedPickFiles.mockResolvedValueOnce(['/file.json'])
+    mockedLoadDataFiles.mockResolvedValueOnce([
+      { data: mockValidData, file: 'file.json' },
+      {
+        data: mockValidData.map((value, ind) => ({ ...value, step: ind })),
+        file: 'file2.json'
+      }
+    ])
+    mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
+    mockedQuickPickUserOrderedValues
+      .mockResolvedValueOnce([
+        {
+          file: 'file.json',
+          key: 'actual'
+        },
+        {
+          file: 'file2.json',
+          key: 'prob'
+        }
+      ])
+      .mockResolvedValueOnce([
+        {
+          file: 'file2.json',
+          key: 'actual'
+        },
+        {
+          file: 'file2.json',
+          key: 'step'
+        }
+      ])
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedQuickPickUserOrderedValues).toHaveBeenCalledTimes(2)
+    expect(result).toStrictEqual(undefined)
+    expect(mockedShowError).toHaveBeenCalledTimes(1)
+    expect(mockedShowError).toHaveBeenCalledWith(
+      'file2.json cannot have more than one metric selected. See [an example](https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#available-configuration-fields) of a plot with multiple x metrics.'
+    )
   })
 
   it('should return early if the user does not pick a title', async () => {
diff --git a/extension/src/pipeline/quickPick.ts b/extension/src/pipeline/quickPick.ts
index b84125b3ee..e62025459d 100644
--- a/extension/src/pipeline/quickPick.ts
+++ b/extension/src/pipeline/quickPick.ts
@@ -26,6 +26,10 @@ export type PlotConfigData = {
 type UnknownValue = Value | ValueTree
 
 type FileFields = { file: string; fields: string[] }[]
+type QuickPickFieldValues = { file: string; field: string }[]
+
+const multipleXMetricsExample =
+  'See [an example](https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#available-configuration-fields) of a plot with multiple x metrics.'
 
 const pickDataFiles = (): Promise<string[] | undefined> =>
   pickFiles(Title.SELECT_PLOT_DATA, {
@@ -33,8 +37,12 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
   })
 
 const formatFieldQuickPickValues = (
-  values: { file: string; field: string }[]
+  values: QuickPickFieldValues | undefined
 ) => {
+  if (!values || values.length === 0) {
+    return
+  }
+
   const formattedFields: PlotConfigDataAxis = {}
 
   for (const { file, field } of values) {
@@ -49,6 +57,77 @@ const formatFieldQuickPickValues = (
   return formattedFields
 }
 
+const verifyFilesHaveSingleField = (files: PlotConfigDataAxis) => {
+  for (const [file, fields] of Object.entries(files)) {
+    if (fields.length > 1) {
+      void Toast.showError(
+        `${file} cannot have more than one metric selected. ${multipleXMetricsExample}`
+      )
+      return
+    }
+  }
+
+  return files
+}
+
+const verifyXFields = (xValues: QuickPickFieldValues | undefined) => {
+  const x = formatFieldQuickPickValues(xValues)
+
+  if (!x) {
+    return
+  }
+
+  const doFilesHaveOneField = verifyFilesHaveSingleField(x)
+
+  if (!doFilesHaveOneField) {
+    return
+  }
+
+  return x
+}
+
+const verifyYFieldsWithMultiXFields = (
+  y: PlotConfigDataAxis,
+  yValues: QuickPickFieldValues,
+  xFieldsLength: number
+) => {
+  if (yValues.length !== xFieldsLength) {
+    void Toast.showError(
+      `Found ${xFieldsLength} x metrics and ${yValues.length} y metric(s). When there are multiple x metrics selected there must be an equal number of y metrics. ${multipleXMetricsExample}`
+    )
+    return
+  }
+
+  const doFilesHaveOneField = verifyFilesHaveSingleField(y)
+
+  if (!doFilesHaveOneField) {
+    return
+  }
+
+  return y
+}
+
+const verifyYFields = (
+  yValues: QuickPickFieldValues | undefined,
+  xFieldsLength = 0
+) => {
+  const y = formatFieldQuickPickValues(yValues)
+
+  if (!y) {
+    return
+  }
+
+  if (xFieldsLength > 1) {
+    return verifyYFieldsWithMultiXFields(
+      y,
+      yValues as QuickPickFieldValues,
+      xFieldsLength
+    )
+  }
+
+  return y
+}
+
 const pickFields = async (
   fileFields: FileFields
 ): Promise<
@@ -74,32 +153,32 @@ const pickFields = async (
 
   const xValues = (await quickPickUserOrderedValues(items, {
     title: Title.SELECT_PLOT_X_METRIC
-  })) as { file: string; field: string }[] | undefined
+  })) as QuickPickFieldValues | undefined
 
-  if (!xValues || xValues.length === 0) {
+  const x = verifyXFields(xValues)
+  if (!x) {
     return
   }
 
   const yValues = (await quickPickUserOrderedValues(
     items.filter(
-      item => item.value === undefined || !xValues.includes(item.value)
+      item => item.value === undefined || !xValues?.includes(item.value)
     ),
     {
       title: Title.SELECT_PLOT_Y_METRIC
     }
   )) as { file: string; field: string }[] | undefined
 
-  if (!yValues || yValues.length === 0) {
+  const y = verifyYFields(yValues, xValues?.length)
+
+  if (!y) {
     return
   }
 
   return {
-    fields: {
-      x: formatFieldQuickPickValues(xValues),
-      y: formatFieldQuickPickValues(yValues)
-    },
-    firstXKey: xValues[0].field,
-    firstYKey: yValues[0].field
+    fields: { x, y },
+    firstXKey: (xValues as QuickPickFieldValues)[0].field,
+    firstYKey: (yValues as QuickPickFieldValues)[0].field
   }
 }