Add error handling for Plot Wizard plots with x fields

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 () => {

extension/src/pipeline/quickPick.ts

@@ -26,15 +26,23 @@ 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, {
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
 
 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
   }
 }