From bdbf500991641fe578fa8744e115ae992420b2ad Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Wed, 25 May 2022 13:30:17 -0400
Subject: [PATCH] Cleanup in comparison table stories

---
 .../comparisonTable/ComparisonTable.test.tsx  | 10 ------
 .../comparisonTable/ComparisonTable.tsx       | 13 ++++---
 .../comparisonTable/ComparisonTableHeader.tsx |  5 ++-
 .../src/stories/ComparisonTable.stories.tsx   | 36 ++++++++++---------
 4 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
index 65673ab223..1e547b4a2f 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
+++ b/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'
 
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.tsx
index 8af03e8720..36938c6da9 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.tsx
+++ b/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>([])
 
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableHeader.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableHeader.tsx
index d780a3f381..15653b2389 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableHeader.tsx
+++ b/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 />
diff --git a/webview/src/stories/ComparisonTable.stories.tsx b/webview/src/stories/ComparisonTable.stories.tsx
index 9c2e83ddee..b3c0b1d78f 100644
--- a/webview/src/stories/ComparisonTable.stories.tsx
+++ b/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,6 +8,7 @@ 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,
@@ -14,26 +16,28 @@ export default {
   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 = (