Add title option to plot wizard (#4786)

julieg18 · web-flow · commit 0226b4c11eda · 2023-10-10T15:49:41.000Z
diff --git 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
@@ -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
@@ -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,13 +10,15 @@ 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)
 const mockedShowError = jest.fn()
 mockedToast.showError = mockedShowError
 
 jest.mock('../fileSystem')
+jest.mock('../vscode/inputBox')
 jest.mock('../vscode/resourcePicker')
 jest.mock('../vscode/quickPick')
 
@@ -250,7 +253,7 @@ 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, x field, y field and title', async () => {
     mockedPickFiles.mockResolvedValueOnce(['/file.json'])
     mockedLoadDataFiles.mockResolvedValueOnce([
       { data: mockValidData, file: '/file.json' }
@@ -259,6 +262,7 @@ describe('pickPlotConfiguration', () => {
     mockedQuickPickValue
       .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
       .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
+    mockedGetInput.mockResolvedValueOnce('Simple Plot')
 
     const result = await pickPlotConfiguration('/')
 
@@ -304,8 +308,13 @@ 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',
       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' }
     })
@@ -392,6 +403,7 @@ describe('pickPlotConfiguration', () => {
       { data: mockValidData, file: 'file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickValue.mockResolvedValueOnce(undefined)
 
     const result = await pickPlotConfiguration('/')
@@ -406,6 +418,7 @@ describe('pickPlotConfiguration', () => {
       { data: mockValidData, file: 'file.json' }
     ])
     mockedQuickPickOne.mockResolvedValueOnce('simple')
+    mockedGetInput.mockResolvedValueOnce('simple_plot')
     mockedQuickPickValue
       .mockResolvedValueOnce('actual')
       .mockResolvedValueOnce(undefined)
@@ -415,4 +428,21 @@ describe('pickPlotConfiguration', () => {
     expect(mockedQuickPickValue).toHaveBeenCalledTimes(2)
     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')
+    mockedQuickPickValue
+      .mockResolvedValueOnce({ file: 'file.json', key: 'actual' })
+      .mockResolvedValueOnce({ file: 'file.json', key: 'prob' })
+    mockedGetInput.mockResolvedValueOnce(undefined)
+
+    const result = await pickPlotConfiguration('/')
+
+    expect(mockedGetInput).toHaveBeenCalledTimes(1)
+    expect(result).toStrictEqual(undefined)
+  })
 })
diff --git 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,9 @@ const pickDataFiles = (): Promise<string[] | undefined> =>
     'Data Formats': ['json', 'csv', 'tsv', 'yaml']
   })
 
-const pickTemplateAndFields = async (
+const pickFields = async (
   fileFields: FileFields
-): Promise<PlotConfigData | undefined> => {
-  const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
-
-  if (!template) {
-    return
-  }
-
+): Promise<Omit<PlotConfigData, 'title' | 'template'> | undefined> => {
   const items = []
 
   for (const { file, fields } of fileFields) {
@@ -65,7 +61,34 @@ const pickTemplateAndFields = async (
     return
   }
 
-  return { template, x, y }
+  return { x, y }
+}
+
+const pickPlotConfigData = async (
+  fileFields: FileFields
+): Promise<PlotConfigData | undefined> => {
+  const template = await quickPickOne(PLOT_TEMPLATES, 'Pick a Plot Template')
+
+  if (!template) {
+    return
+  }
+
+  const fields = await pickFields(fileFields)
+
+  if (!fields) {
+    return
+  }
+
+  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[]) => {
@@ -222,11 +245,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
@@ -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
@@ -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',
