Add select column options to table header cell context menu (#4295)

extension/src/experiments/index.ts

@@ -598,6 +598,7 @@ export class Experiments extends BaseRepository<TableData> {
       () => this.getWebview(),
       () => this.notifyChanged(),
       () => this.selectColumns(),
+      () => this.selectFirstColumns(),
       (branchesSelected: string[]) => this.selectBranches(branchesSelected),
       () => this.data.update()
     )

extension/src/experiments/webview/messages.ts

@@ -33,6 +33,7 @@ export class WebviewMessages {
   private readonly getWebview: () => BaseWebview<TableData> | undefined
   private readonly notifyChanged: () => void
   private readonly selectColumns: () => Promise<void>
+  private readonly selectFirstColumns: () => Promise<void>
 
   private readonly selectBranches: (
     branchesSelected: string[]
@@ -48,6 +49,7 @@ export class WebviewMessages {
     getWebview: () => BaseWebview<TableData> | undefined,
     notifyChanged: () => void,
     selectColumns: () => Promise<void>,
+    selectFirstColumns: () => Promise<void>,
     selectBranches: (
       branchesSelected: string[]
     ) => Promise<string[] | undefined>,
@@ -60,6 +62,7 @@ export class WebviewMessages {
     this.getWebview = getWebview
     this.notifyChanged = notifyChanged
     this.selectColumns = selectColumns
+    this.selectFirstColumns = selectFirstColumns
     this.selectBranches = selectBranches
     this.update = update
   }
@@ -131,6 +134,8 @@ export class WebviewMessages {
 
       case MessageFromWebviewType.SELECT_COLUMNS:
         return this.setColumnsStatus()
+      case MessageFromWebviewType.SELECT_FIRST_COLUMNS:
+        return this.setFirstColumns()
 
       case MessageFromWebviewType.FOCUS_FILTERS_TREE:
         return this.focusFiltersTree()
@@ -322,6 +327,15 @@ export class WebviewMessages {
     )
   }
 
+  private setFirstColumns() {
+    void this.selectFirstColumns()
+    sendTelemetryEvent(
+      EventName.VIEWS_EXPERIMENTS_TABLE_SELECT_COLUMNS,
+      undefined,
+      undefined
+    )
+  }
+
   private addColumnSort(sort: SortDefinition) {
     this.experiments.addSort(sort)
     sendTelemetryEvent(

extension/src/telemetry/constants.ts

@@ -57,6 +57,9 @@ export const EventName = Object.assign(
       'views.experimentsTable.selectColumns',
     VIEWS_EXPERIMENTS_TABLE_SELECT_EXPERIMENTS_FOR_PLOTS:
       'views.experimentsTable.selectExperimentsForPlots',
+    VIEWS_EXPERIMENTS_TABLE_SELECT_FIRST_COLUMNS:
+      'views.experimentsTable.selectFirstColumns',
+
     VIEWS_EXPERIMENTS_TABLE_SET_MAX_HEADER_HEIGHT:
       'views.experimentsTable.updateHeaderMaxHeight',
     VIEWS_EXPERIMENTS_TABLE_SHOW_LESS_COMMITS:
@@ -251,6 +254,7 @@ export interface IEventNamePropertyMapping {
   [EventName.VIEWS_EXPERIMENTS_TABLE_SELECT_EXPERIMENTS_FOR_PLOTS]: {
     experimentCount: number
   }
+  [EventName.VIEWS_EXPERIMENTS_TABLE_SELECT_FIRST_COLUMNS]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_SHOW_MORE_COMMITS]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_SHOW_LESS_COMMITS]: undefined
   [EventName.VIEWS_EXPERIMENTS_TABLE_OPEN_PARAMS_FILE]: {

extension/src/test/suite/experiments/index.test.ts

@@ -1018,6 +1018,38 @@ suite('Experiments Test Suite', () => {
       expect(messageSpy).to.be.calledWithMatch(allColumnsUnselected)
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
+    it('should be able to handle a message to select the first columns', async () => {
+      const { experiments, messageSpy } = setupExperimentsAndMockCommands()
+
+      const webview = await experiments.showWebview()
+      messageSpy.resetHistory()
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      const movedColumn = 'metrics:summary.json:val_accuracy'
+
+      const mockShowQuickPick = stub(window, 'showQuickPick') as SinonStub<
+        [items: readonly QuickPickItem[], options: QuickPickOptionsWithTitle],
+        Thenable<QuickPickItemWithValue<{ path: string }>[] | undefined>
+      >
+      mockShowQuickPick.resolves([
+        {
+          value: { path: movedColumn },
+          label: movedColumn
+        }
+      ])
+
+      const tableChangePromise = experimentsUpdatedEvent(experiments)
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.SELECT_FIRST_COLUMNS
+      })
+      await tableChangePromise
+
+      const [id, firstColumn] = messageSpy.lastCall.args[0].columnOrder
+
+      expect(id).to.equal('id')
+      expect(firstColumn).to.equal(movedColumn)
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
     it('should be able to handle a message to focus the sorts tree', async () => {
       const { experiments } = buildExperiments({ disposer: disposable })
 

extension/src/webview/contract.ts

@@ -51,6 +51,7 @@ export enum MessageFromWebviewType {
   HIDE_EXPERIMENTS_TABLE_COLUMN = 'hide-experiments-table-column',
   SELECT_EXPERIMENTS = 'select-experiments',
   SELECT_COLUMNS = 'select-columns',
+  SELECT_FIRST_COLUMNS = 'select-first-columns',
   SELECT_PLOTS = 'select-plots',
   SET_EXPERIMENTS_FOR_PLOTS = 'set-experiments-for-plots',
   SET_EXPERIMENTS_AND_OPEN_PLOTS = 'set-experiments-and-open-plots',
@@ -216,6 +217,7 @@ export type MessageFromWebview =
   | { type: MessageFromWebviewType.REFRESH_EXP_DATA }
   | { type: MessageFromWebviewType.REFRESH_REVISIONS }
   | { type: MessageFromWebviewType.SELECT_COLUMNS }
+  | { type: MessageFromWebviewType.SELECT_FIRST_COLUMNS }
   | { type: MessageFromWebviewType.FOCUS_FILTERS_TREE }
   | { type: MessageFromWebviewType.FOCUS_SORTS_TREE }
   | { type: MessageFromWebviewType.OPEN_PLOTS_WEBVIEW }

webview/src/experiments/components/App.test.tsx

@@ -737,10 +737,10 @@ describe('App', () => {
       advanceTimersByTime(100)
 
       const menuitems = screen.getAllByRole('menuitem')
-      expect(menuitems).toHaveLength(6)
+      expect(menuitems).toHaveLength(8)
       expect(
         menuitems.filter(item => !item.className.includes('disabled'))
-      ).toHaveLength(3)
+      ).toHaveLength(5)
 
       fireEvent.keyDown(paramsFileHeader, { bubbles: true, key: 'Escape' })
       expect(screen.queryAllByRole('menuitem')).toHaveLength(0)
@@ -761,7 +761,7 @@ describe('App', () => {
       expect(disabledMenuItem).toBeDefined()
 
       disabledMenuItem && fireEvent.click(disabledMenuItem, { bubbles: true })
-      expect(screen.queryAllByRole('menuitem')).toHaveLength(6)
+      expect(screen.queryAllByRole('menuitem')).toHaveLength(8)
     })
 
     it('should have the same enabled options in the empty placeholders', () => {
@@ -783,6 +783,8 @@ describe('App', () => {
         expect(menuitems).toStrictEqual([
           'Hide Column',
           'Set Max Header Height',
+          'Select Columns',
+          'Select First Columns',
           'Sort Ascending',
           'Sort Descending'
         ])
@@ -807,12 +809,54 @@ describe('App', () => {
           .filter(item => !item.className.includes('disabled'))
           .map(item => item.textContent)
 
-        expect(menuitems).toStrictEqual(['Set Max Header Height'])
+        expect(menuitems).toStrictEqual([
+          'Set Max Header Height',
+          'Select Columns',
+          'Select First Columns'
+        ])
 
         fireEvent.keyDown(segment, { bubbles: true, key: 'Escape' })
       }
     })
 
+    it('should send the correct message when Select Columns is clicked', () => {
+      renderTableWithPlaceholder()
+      const placeholders = screen.getAllByTestId(/header-Created/)
+      const placeholder = placeholders[0]
+      fireEvent.contextMenu(placeholder, { bubbles: true })
+      advanceTimersByTime(100)
+
+      const selectOption = screen.getByText('Select Columns')
+
+      mockPostMessage.mockClear()
+
+      fireEvent.click(selectOption)
+
+      expect(mockPostMessage).toHaveBeenCalledTimes(1)
+      expect(mockPostMessage).toHaveBeenCalledWith({
+        type: MessageFromWebviewType.SELECT_COLUMNS
+      })
+    })
+
+    it('should send the correct message when Select First Columns is clicked', () => {
+      renderTableWithPlaceholder()
+      const placeholders = screen.getAllByTestId(/header-Created/)
+      const placeholder = placeholders[0]
+      fireEvent.contextMenu(placeholder, { bubbles: true })
+      advanceTimersByTime(100)
+
+      const selectOption = screen.getByText('Select First Columns')
+
+      mockPostMessage.mockClear()
+
+      fireEvent.click(selectOption)
+
+      expect(mockPostMessage).toHaveBeenCalledTimes(1)
+      expect(mockPostMessage).toHaveBeenCalledWith({
+        type: MessageFromWebviewType.SELECT_FIRST_COLUMNS
+      })
+    })
+
     describe('Hiding a column from its empty placeholder', () => {
       it('should send the column id and not the placeholder id as the message payload', () => {
         renderTableWithPlaceholder()

webview/src/experiments/components/table/header/ContextMenuContent.tsx

@@ -114,11 +114,26 @@ export const getMenuOptions = (
       }
     },
     {
+      divider: true,
       id: 'update-header-depth',
       label: 'Set Max Header Height',
       message: {
         type: MessageFromWebviewType.SET_EXPERIMENTS_HEADER_HEIGHT
       }
+    },
+    {
+      id: 'select-columns',
+      label: 'Select Columns',
+      message: {
+        type: MessageFromWebviewType.SELECT_COLUMNS
+      }
+    },
+    {
+      id: 'select-first-columns',
+      label: 'Select First Columns',
+      message: {
+        type: MessageFromWebviewType.SELECT_FIRST_COLUMNS
+      }
     }
   ]