Cleanup in comparison table stories (#1771)

Author: sroy3

webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx

@@ -213,16 +213,6 @@ describe('ComparisonTable', () => {
     expect(headers).toStrictEqual([...namedRevisions, newRevName])
   })
 
-  it('should pin the current pinned column on first render', () => {
-    const pinnedRevision = 'main'
-
-    renderTable({ ...basicProps, currentPinnedColumn: pinnedRevision })
-
-    const [pinnedColumn] = getHeaders()
-
-    expect(pinnedColumn.textContent).toBe(pinnedRevision)
-  })
-
   it('should display a refresh button for each revision that has a missing image', () => {
     const revisionWithNoData = 'missing-data'
 

webview/src/plots/components/comparisonTable/ComparisonTable.tsx

@@ -15,17 +15,16 @@ import plotsStyles from '../styles.module.scss'
 import { withScale } from '../../../util/styles'
 import { sendMessage } from '../../../shared/vscode'
 
-export interface ComparisonTableProps
-  extends Omit<PlotsComparisonData, 'sectionName' | 'size'> {
-  currentPinnedColumn?: string
-}
+export type ComparisonTableProps = Omit<
+  PlotsComparisonData,
+  'sectionName' | 'size'
+>
 
 export const ComparisonTable: React.FC<ComparisonTableProps> = ({
   plots,
-  revisions,
-  currentPinnedColumn
+  revisions
 }) => {
-  const pinnedColumn = useRef(currentPinnedColumn || '')
+  const pinnedColumn = useRef('')
   const [columns, setColumns] = useState<ComparisonTableColumn[]>([])
   const [comparisonPlots, setComparisonPlots] = useState<ComparisonPlots>([])
 

webview/src/plots/components/comparisonTable/ComparisonTableHeader.tsx

@@ -21,7 +21,10 @@ export const ComparisonTableHeader: React.FC<ComparisonTableHeaderProps> = ({
   })
 
   return (
-    <div className={styles.header}>
+    <div
+      className={styles.header}
+      data-testid={`${children?.toString().split(',')[0]}-header`}
+    >
       {!isPinned && <GripIcon className={styles.gripIcon} />}
       <button className={pinClasses} onClick={onClicked}>
         <Pin />

webview/src/stories/ComparisonTable.stories.tsx

@@ -1,4 +1,5 @@
 import { Meta, Story } from '@storybook/react/types-6-0'
+import { fireEvent, within } from '@testing-library/react'
 import React from 'react'
 import { ComparisonRevisionData } from 'dvc/src/plots/webview/contract'
 import comparisonTableFixture from 'dvc/src/test/fixtures/plotsDiff/comparison'
@@ -7,33 +8,36 @@ import {
   ComparisonTableProps
 } from '../plots/components/comparisonTable/ComparisonTable'
 import { Theme } from '../shared/components/theme/Theme'
+import { DragDropProvider } from '../shared/components/dragDrop/DragDropContext'
 
 export default {
   args: comparisonTableFixture,
   component: ComparisonTable,
   title: 'Comparison Table'
 } as Meta
 
-const Template: Story<ComparisonTableProps> = ({
-  plots,
-  revisions,
-  currentPinnedColumn
-}) => (
-  <Theme>
-    <ComparisonTable
-      plots={plots}
-      revisions={revisions}
-      currentPinnedColumn={currentPinnedColumn}
-    />
-  </Theme>
-)
+const Template: Story<ComparisonTableProps> = ({ plots, revisions }) => {
+  return (
+    <Theme>
+      <DragDropProvider>
+        <ComparisonTable plots={plots} revisions={revisions} />
+      </DragDropProvider>
+    </Theme>
+  )
+}
 
 export const Basic = Template.bind({})
 
 export const WithPinnedColumn = Template.bind({})
-WithPinnedColumn.args = {
-  ...comparisonTableFixture,
-  currentPinnedColumn: 'main'
+WithPinnedColumn.parameters = {
+  chromatic: { delay: 300 }
+}
+WithPinnedColumn.play = async ({ canvasElement }) => {
+  const canvas = within(canvasElement)
+  const mainHeader = await canvas.findByTestId('main-header')
+  const pin = within(mainHeader).getByRole('button')
+
+  fireEvent.click(pin)
 }
 
 const removeSingleImage = (