From 5f5babf9308eac531850c1896c988f0fc78e3ca1 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Tue, 19 Jul 2022 13:18:15 -0400
Subject: [PATCH 1/8] Change the drag and drop feedback for the comparison
 table

---
 .../components/table/TableHeader.tsx          |  5 +-
 .../checkpointPlots/CheckpointPlots.tsx       |  5 +-
 .../comparisonTable/ComparisonTableHead.tsx   |  6 +-
 .../components/comparisonTable/DropTarget.tsx | 13 +---
 .../comparisonTable/styles.module.scss        |  6 ++
 .../templatePlots/TemplatePlotsGrid.tsx       |  5 +-
 .../components/dragDrop/DragDropContainer.tsx | 68 +++++++++++++------
 .../components/dragDrop/DragDropWorkbench.tsx |  4 +-
 8 files changed, 64 insertions(+), 48 deletions(-)

diff --git a/webview/src/experiments/components/table/TableHeader.tsx b/webview/src/experiments/components/table/TableHeader.tsx
index d048596a08..e11402461c 100644
--- a/webview/src/experiments/components/table/TableHeader.tsx
+++ b/webview/src/experiments/components/table/TableHeader.tsx
@@ -58,10 +58,7 @@ export const ColumnDragHandle: React.FC<{
         id={column.id}
         disabled={disabled}
         group={'experiment-table'}
-        dropTarget={{
-          element: DropTarget,
-          wrapperTag: 'div'
-        }}
+        dropTarget={DropTarget}
         onDragOver={onDragOver}
         onDragStart={onDragStart}
         onDrop={onDrop}
diff --git a/webview/src/plots/components/checkpointPlots/CheckpointPlots.tsx b/webview/src/plots/components/checkpointPlots/CheckpointPlots.tsx
index 5e9c47cec0..9a0493ec07 100644
--- a/webview/src/plots/components/checkpointPlots/CheckpointPlots.tsx
+++ b/webview/src/plots/components/checkpointPlots/CheckpointPlots.tsx
@@ -63,10 +63,7 @@ export const CheckpointPlots: React.FC<CheckpointPlotsProps> = ({
         disabledDropIds={[]}
         items={items as JSX.Element[]}
         group="live-plots"
-        dropTarget={{
-          element: <DropTarget />,
-          wrapperTag: 'div'
-        }}
+        dropTarget={<DropTarget />}
         wrapperComponent={
           useVirtualizedGrid
             ? {
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx
index 0db1018580..aeb9c09513 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx
@@ -52,10 +52,8 @@ export const ComparisonTableHead: React.FC<ComparisonTableHeadProps> = ({
           disabledDropIds={[pinnedColumn]}
           items={items}
           group="comparison"
-          dropTarget={{
-            element: <DropTarget />,
-            wrapperTag: 'th'
-          }}
+          dropTarget={<DropTarget />}
+          shouldShowOnDrag
         />
       </tr>
     </thead>
diff --git a/webview/src/plots/components/comparisonTable/DropTarget.tsx b/webview/src/plots/components/comparisonTable/DropTarget.tsx
index f34781a27d..b0f11c7c02 100644
--- a/webview/src/plots/components/comparisonTable/DropTarget.tsx
+++ b/webview/src/plots/components/comparisonTable/DropTarget.tsx
@@ -1,15 +1,6 @@
 import React from 'react'
-import { Icon } from '../../../shared/components/Icon'
-import { Ellipsis } from '../../../shared/components/icons'
-import styles from '../styles.module.scss'
+import styles from './styles.module.scss'
 
 export const DropTarget: React.FC = () => (
-  <div className={styles.dropTarget} data-testid="comparison-drop-target">
-    <Icon
-      icon={Ellipsis}
-      className={styles.smallDropIcon}
-      width={15}
-      height={15}
-    />
-  </div>
+  <div className={styles.dropTarget} data-testid="comparison-drop-target"></div>
 )
diff --git a/webview/src/plots/components/comparisonTable/styles.module.scss b/webview/src/plots/components/comparisonTable/styles.module.scss
index 8b67d7268e..95235ef63e 100644
--- a/webview/src/plots/components/comparisonTable/styles.module.scss
+++ b/webview/src/plots/components/comparisonTable/styles.module.scss
@@ -157,3 +157,9 @@ $gap: 4px;
 .experimentName {
   color: $meta-cell-color;
 }
+
+.dropTarget {
+  width: 2px;
+  height: 100%;
+  border-right: 2px dashed $accent-color;
+}
diff --git a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
index 645c97acbe..6559e68b95 100644
--- a/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
+++ b/webview/src/plots/components/templatePlots/TemplatePlotsGrid.tsx
@@ -92,10 +92,7 @@ export const TemplatePlotsGrid: React.FC<TemplatePlotsGridProps> = ({
       items={items as JSX.Element[]}
       group={groupId}
       onDrop={onDropInSection}
-      dropTarget={{
-        element: <DropTarget />,
-        wrapperTag: 'div'
-      }}
+      dropTarget={<DropTarget />}
       wrapperComponent={
         useVirtualizedGrid
           ? {
diff --git a/webview/src/shared/components/dragDrop/DragDropContainer.tsx b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
index 0ba5ed3960..2f0597dad3 100644
--- a/webview/src/shared/components/dragDrop/DragDropContainer.tsx
+++ b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
@@ -1,3 +1,4 @@
+import cx from 'classnames'
 import React, {
   DragEvent,
   useEffect,
@@ -8,6 +9,7 @@ import React, {
 import { useDispatch, useSelector } from 'react-redux'
 import { DragEnterDirection, getDragEnterDirection } from './util'
 import { changeRef } from './dragDropSlice'
+import styles from './styles.module.scss'
 import { getIDIndex, getIDWithoutIndex } from '../../../util/ids'
 import { Any } from '../../../util/objects'
 import { RootState } from '../../../plots/store'
@@ -34,26 +36,23 @@ export type OnDrop = (
   position: number
 ) => void
 
-export type DropTargetInfo = {
-  element: JSX.Element
-  wrapperTag: 'div' | 'th'
-}
-
 export const makeTarget = (
-  dropTarget: DropTargetInfo,
+  dropTarget: JSX.Element,
   handleDragOver: DragEventHandler<HTMLElement>,
   handleOnDrop: DragEventHandler<HTMLElement>,
-  id: string
+  id: string,
+  className?: string
 ) => (
-  <dropTarget.wrapperTag
+  <div
     data-testid="drop-target"
     key="drop-target"
     onDragOver={handleDragOver}
     onDrop={handleOnDrop}
     id={`${id}__drop`}
+    className={className}
   >
-    {dropTarget.element}
-  </dropTarget.wrapperTag>
+    {dropTarget}
+  </div>
 )
 interface DragDropContainerProps {
   order: string[]
@@ -62,13 +61,14 @@ interface DragDropContainerProps {
   items: JSX.Element[] // Every item must have a id prop for drag and drop to work
   group: string
   onDrop?: OnDrop
-  dropTarget: DropTargetInfo
+  dropTarget: JSX.Element
   wrapperComponent?: {
     component: React.FC<WrapperProps>
     props: {
       [key: string]: Any
     }
   }
+  shouldShowOnDrag?: boolean
 }
 
 export const DragDropContainer: React.FC<DragDropContainerProps> = ({
@@ -79,7 +79,8 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
   group,
   onDrop,
   dropTarget,
-  wrapperComponent
+  wrapperComponent,
+  shouldShowOnDrag
   // eslint-disable-next-line sonarjs/cognitive-complexity
 }) => {
   const [draggedOverId, setDraggedOverId] = useState('')
@@ -166,7 +167,9 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
     const orderIdxChange = orderIdxTune(direction, droppedIndex > draggedIndex)
     const orderIdxChanged = droppedIndex + orderIdxChange
     const isEnabled = !disabledDropIds.includes(order[orderIdxChanged])
-    setDraggedId('')
+    draggedOverIdTimeout.current = window.setTimeout(() => {
+      setDraggedId('')
+    }, 0)
 
     if (isEnabled && isSameGroup(e.dataTransfer.getData('group'), group)) {
       applyDrop(e, orderIdxChanged, draggedIndex, newOrder, oldDraggedId)
@@ -207,20 +210,47 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
       onDragEnter={handleDragEnter}
       onDrop={handleOnDrop}
       draggable={!disabledDropIds.includes(id)}
-      style={(id === draggedId && { display: 'none' }) || draggable.props.style}
+      className={cx(draggable.props.className, {
+        [styles.hiddenOnDrag]: !shouldShowOnDrag && id === draggedId
+      })}
     />
   )
 
-  const wrappedItems = items.flatMap(draggable => {
+  const wrappedItems = items.map(draggable => {
     const { id } = draggable.props
     const item = id && buildItem(id, draggable)
 
     if (id === draggedOverId) {
-      const target = makeTarget(dropTarget, handleDragOver, handleOnDrop, id)
+      const targetClassName = shouldShowOnDrag
+        ? cx(styles.dropTargetWhenShowingOnDrag, {
+            [styles.dropTargetWhenShowingOnDragLeft]:
+              direction === DragEnterDirection.LEFT,
+            [styles.dropTargetWhenShowingOnDragRight]:
+              direction === DragEnterDirection.RIGHT
+          })
+        : undefined
+      const target = makeTarget(
+        dropTarget,
+        handleDragOver,
+        handleOnDrop,
+        id,
+        targetClassName
+      )
+      const Tag = item.type
+      const itemWithTag = shouldShowOnDrag ? <div {...item.props} /> : item
+      const block =
+        direction === DragEnterDirection.RIGHT
+          ? [itemWithTag, target]
+          : [target, itemWithTag]
 
-      return direction === DragEnterDirection.RIGHT
-        ? [item, target]
-        : [target, item]
+      if (shouldShowOnDrag) {
+        return (
+          <Tag key={draggable.key} className={styles.newBlockWhenShowingOnDrag}>
+            {block}
+          </Tag>
+        )
+      }
+      return block
     }
 
     return item
diff --git a/webview/src/shared/components/dragDrop/DragDropWorkbench.tsx b/webview/src/shared/components/dragDrop/DragDropWorkbench.tsx
index cfaa3d4d7f..fef1886459 100644
--- a/webview/src/shared/components/dragDrop/DragDropWorkbench.tsx
+++ b/webview/src/shared/components/dragDrop/DragDropWorkbench.tsx
@@ -1,5 +1,5 @@
 import React, { DragEvent, useContext } from 'react'
-import { DropTargetInfo, makeTarget } from './DragDropContainer'
+import { makeTarget } from './DragDropContainer'
 import { DragDropContext, DragDropContextValue } from './DragDropContext'
 
 export type OnDrop = (draggedId: string, draggedOverId: string) => void
@@ -10,7 +10,7 @@ export interface DraggableProps {
   id: string
   group: string
   disabled: boolean
-  dropTarget: DropTargetInfo
+  dropTarget: JSX.Element
   children: JSX.Element
   onDrop?: OnDrop
   onDragStart?: OnDragStart

From d027b40695fee54b018136431fdcacc1d92f55b0 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Wed, 20 Jul 2022 16:12:52 -0400
Subject: [PATCH 2/8] Fix tests

---
 .../src/experiments/components/App.test.tsx   |  7 ++++-
 .../components/table/Table.test.tsx           | 16 +++++++---
 webview/src/plots/components/App.test.tsx     | 10 +++----
 .../comparisonTable/ComparisonTable.test.tsx  |  2 +-
 .../components/dragDrop/DragDropContainer.tsx | 28 +++++++++++-------
 .../components/dragDrop/styles.module.scss    | 22 ++++++++++++++
 webview/src/test/dragDrop.ts                  | 29 ++++++++++++++-----
 7 files changed, 86 insertions(+), 28 deletions(-)
 create mode 100644 webview/src/shared/components/dragDrop/styles.module.scss

diff --git a/webview/src/experiments/components/App.test.tsx b/webview/src/experiments/components/App.test.tsx
index 2bf82d9ccc..5fe65552dd 100644
--- a/webview/src/experiments/components/App.test.tsx
+++ b/webview/src/experiments/components/App.test.tsx
@@ -182,7 +182,12 @@ describe('App', () => {
     const headerB = screen.getByText('B')
     const headerD = screen.getByText('D')
 
-    dragAndDrop(headerB, headerD, DragEnterDirection.AUTO)
+    dragAndDrop(
+      headerB,
+      // eslint-disable-next-line testing-library/no-node-access
+      headerD.parentElement?.parentElement || headerD,
+      DragEnterDirection.AUTO
+    )
 
     await expectHeaders(['A', 'C', 'D', 'B'])
   })
diff --git a/webview/src/experiments/components/table/Table.test.tsx b/webview/src/experiments/components/table/Table.test.tsx
index b04b9a000f..bc840c43ac 100644
--- a/webview/src/experiments/components/table/Table.test.tsx
+++ b/webview/src/experiments/components/table/Table.test.tsx
@@ -273,7 +273,9 @@ describe('Table', () => {
 
       dragAndDrop(
         screen.getByText('B'),
-        screen.getByText('C'),
+        // eslint-disable-next-line testing-library/no-node-access
+        screen.getByText('C').parentElement?.parentElement ||
+          screen.getByText('C'),
         DragEnterDirection.AUTO
       )
 
@@ -281,7 +283,9 @@ describe('Table', () => {
 
       dragAndDrop(
         screen.getByText('A'),
-        screen.getByText('B'),
+        // eslint-disable-next-line testing-library/no-node-access
+        screen.getByText('B').parentElement?.parentElement ||
+          screen.getByText('B'),
         DragEnterDirection.AUTO
       )
 
@@ -370,7 +374,9 @@ describe('Table', () => {
 
       dragAndDrop(
         screen.getByText('process'),
-        screen.getByText('loss'),
+        // eslint-disable-next-line testing-library/no-node-access
+        screen.getByText('loss').parentElement?.parentElement ||
+          screen.getByText('loss'),
         DragEnterDirection.AUTO
       )
 
@@ -382,7 +388,9 @@ describe('Table', () => {
 
       dragAndDrop(
         screen.getByText('summary.json'),
-        screen.getByText('test'),
+        // eslint-disable-next-line testing-library/no-node-access
+        screen.getByText('test').parentElement?.parentElement ||
+          screen.getByText('test'),
         DragEnterDirection.AUTO
       )
 
diff --git a/webview/src/plots/components/App.test.tsx b/webview/src/plots/components/App.test.tsx
index 9328a069b6..4bde293e7c 100644
--- a/webview/src/plots/components/App.test.tsx
+++ b/webview/src/plots/components/App.test.tsx
@@ -742,7 +742,7 @@ describe('App', () => {
 
     dragEnter(
       anotherSingleViewPlot,
-      movedSingleViewPlot,
+      CSS.escape(movedSingleViewPlot.id),
       DragEnterDirection.LEFT
     )
 
@@ -793,7 +793,7 @@ describe('App', () => {
     const headers = screen.getAllByRole('columnheader')
     const bottomSection = screen.getByTestId(NewSectionBlock.BOTTOM)
 
-    dragEnter(headers[1], bottomSection, DragEnterDirection.LEFT)
+    dragEnter(headers[1], bottomSection.id, DragEnterDirection.LEFT)
 
     const bottomDropIcon = screen.queryByTestId(
       `${NewSectionBlock.BOTTOM}_drop-icon`
@@ -825,7 +825,7 @@ describe('App', () => {
 
     const plots = screen.getAllByTestId(/^plot_/)
 
-    dragEnter(plots[1], plots[0], DragEnterDirection.LEFT)
+    dragEnter(plots[1], CSS.escape(plots[0].id), DragEnterDirection.LEFT)
 
     const plotsWithDropTarget = screen.getAllByTestId(/^plot_/)
     expect(plotsWithDropTarget.map(plot => plot.id)).toStrictEqual([
@@ -842,7 +842,7 @@ describe('App', () => {
     })
 
     const plots = screen.getAllByTestId(/^plot_/)
-    dragEnter(plots[0], plots[1], DragEnterDirection.RIGHT)
+    dragEnter(plots[0], CSS.escape(plots[1].id), DragEnterDirection.RIGHT)
 
     const plotsWithDropTarget = screen.getAllByTestId(/^plot_/)
 
@@ -862,7 +862,7 @@ describe('App', () => {
     const plots = screen.getAllByTestId(/^plot_/)
     expect(plots[1].style.display).not.toBe('none')
 
-    dragEnter(plots[1], plots[1], DragEnterDirection.LEFT)
+    dragEnter(plots[1], CSS.escape(plots[1].id), DragEnterDirection.LEFT)
 
     expect(plots[1].style.display).toBe('none')
   })
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
index 1e7613ab07..3c410a2b93 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
@@ -380,7 +380,7 @@ describe('ComparisonTable', () => {
       pinSecondColumn()
       const [endingNode, startingNode] = getHeaders()
 
-      dragEnter(startingNode, endingNode, DragEnterDirection.LEFT)
+      dragEnter(startingNode, endingNode.id, DragEnterDirection.LEFT)
 
       const headers = getHeaders()
 
diff --git a/webview/src/shared/components/dragDrop/DragDropContainer.tsx b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
index 2f0597dad3..82d137ea4a 100644
--- a/webview/src/shared/components/dragDrop/DragDropContainer.tsx
+++ b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
@@ -107,8 +107,8 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
   const handleDragStart = (e: DragEvent<HTMLElement>) => {
     const { id } = e.currentTarget
     const idx = order.indexOf(id)
-    let toIdx = idx + 1
-    if (toIdx === order.length) {
+    let toIdx = shouldShowOnDrag ? idx : idx + 1
+    if (!shouldShowOnDrag && toIdx === order.length) {
       toIdx = idx - 1
 
       if (toIdx === -1) {
@@ -153,6 +153,12 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
   }
 
   const handleOnDrop = (e: DragEvent<HTMLElement>) => {
+    draggedOverIdTimeout.current = window.setTimeout(() => {
+      setDraggedId('')
+    }, 0)
+    if (e.dataTransfer.getData('itemId') === draggedOverId) {
+      return
+    }
     const newOrder = [...order]
     const oldDraggedId = e.dataTransfer.getData('itemId')
     const isNew = !order.includes(draggedId)
@@ -167,9 +173,6 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
     const orderIdxChange = orderIdxTune(direction, droppedIndex > draggedIndex)
     const orderIdxChanged = droppedIndex + orderIdxChange
     const isEnabled = !disabledDropIds.includes(order[orderIdxChanged])
-    draggedOverIdTimeout.current = window.setTimeout(() => {
-      setDraggedId('')
-    }, 0)
 
     if (isEnabled && isSameGroup(e.dataTransfer.getData('group'), group)) {
       applyDrop(e, orderIdxChanged, draggedIndex, newOrder, oldDraggedId)
@@ -210,13 +213,14 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
       onDragEnter={handleDragEnter}
       onDrop={handleOnDrop}
       draggable={!disabledDropIds.includes(id)}
-      className={cx(draggable.props.className, {
-        [styles.hiddenOnDrag]: !shouldShowOnDrag && id === draggedId
-      })}
+      style={
+        (!shouldShowOnDrag && id === draggedId && { display: 'none' }) ||
+        draggable.props.style
+      }
     />
   )
 
-  const wrappedItems = items.map(draggable => {
+  const wrappedItems = items.flatMap(draggable => {
     const { id } = draggable.props
     const item = id && buildItem(id, draggable)
 
@@ -237,7 +241,11 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
         targetClassName
       )
       const Tag = item.type
-      const itemWithTag = shouldShowOnDrag ? <div {...item.props} /> : item
+      const itemWithTag = shouldShowOnDrag ? (
+        <div key="item" {...item.props} />
+      ) : (
+        item
+      )
       const block =
         direction === DragEnterDirection.RIGHT
           ? [itemWithTag, target]
diff --git a/webview/src/shared/components/dragDrop/styles.module.scss b/webview/src/shared/components/dragDrop/styles.module.scss
new file mode 100644
index 0000000000..878b7e6d5d
--- /dev/null
+++ b/webview/src/shared/components/dragDrop/styles.module.scss
@@ -0,0 +1,22 @@
+.newBlockWhenShowingOnDrag {
+  position: relative;
+}
+
+.dropTargetWhenShowingOnDrag {
+  position: absolute;
+  top: 0;
+  height: 100%;
+  z-index: 2;
+}
+
+.dropTargetWhenShowingOnDragLeft {
+  left: -2px;
+}
+
+.dropTargetWhenShowingOnDragRight {
+  right: -2px;
+}
+
+.hiddenOnDrag {
+  display: none;
+}
diff --git a/webview/src/test/dragDrop.ts b/webview/src/test/dragDrop.ts
index 401058c9ec..64750dc0d0 100644
--- a/webview/src/test/dragDrop.ts
+++ b/webview/src/test/dragDrop.ts
@@ -19,7 +19,7 @@ export const createBubbledEvent = (type: string, props = {}) => {
 
 export const dragEnter = (
   dragged: HTMLElement,
-  draggedOver: HTMLElement,
+  draggedOverId: string,
   direction: DragDropUtils.DragEnterDirection
 ) => {
   jest.useFakeTimers()
@@ -31,9 +31,14 @@ export const dragEnter = (
     jest.advanceTimersByTime(1)
   })
 
+  let draggedOver =
+    (draggedOverId && document.querySelector(`#${draggedOverId}`)) || null
+
   act(() => {
-    draggedOver.dispatchEvent(createBubbledEvent('dragenter'))
+    draggedOver?.dispatchEvent(createBubbledEvent('dragenter'))
   })
+  draggedOver =
+    (draggedOverId && document.querySelector(`#${draggedOverId}`)) || null
 
   act(() => {
     if (direction !== DragDropUtils.DragEnterDirection.AUTO) {
@@ -45,9 +50,10 @@ export const dragEnter = (
       jest
         .spyOn(DragDropUtils, 'getEventCurrentTargetDistances')
         .mockImplementationOnce(() => ({ left, right }))
-      draggedOver.dispatchEvent(dragOverEvent)
+      draggedOver?.dispatchEvent(dragOverEvent)
+      jest.advanceTimersByTime(1)
     } else {
-      draggedOver.dispatchEvent(createBubbledEvent('dragover'))
+      draggedOver?.dispatchEvent(createBubbledEvent('dragover'))
     }
   })
 
@@ -59,19 +65,28 @@ export const dragAndDrop = (
   direction: DragDropUtils.DragEnterDirection = DragDropUtils.DragEnterDirection
     .LEFT
 ) => {
-  dragEnter(startingNode, endingNode, direction)
+  // When showing element on drag, the dragged over element is being recreacted to be wrapped in another element, thus the endingNode does not exist as is in the document
+  // CSS.escape is needed for weird ids in the experiments table (ex.: :loss)
+  const endingNodeId = CSS.escape(endingNode.id)
+  dragEnter(startingNode, endingNodeId, direction)
 
   jest.useFakeTimers()
   act(() => {
-    endingNode.dispatchEvent(createBubbledEvent('drop'))
+    jest.advanceTimersByTime(1)
+  })
+  act(() => {
+    const endingNodeWithId =
+      (endingNodeId && document.querySelector(`#${endingNodeId}`)) || null
+    endingNodeWithId?.dispatchEvent(createBubbledEvent('drop'))
   })
 
   act(() => {
     jest.advanceTimersByTime(1)
   })
-  jest.useRealTimers()
 
   act(() => {
     startingNode.dispatchEvent(createBubbledEvent('dragend'))
   })
+
+  jest.useRealTimers()
 }

From 282451e76084ceddb012da0cdf5f75b26f99c6f5 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 21 Jul 2022 10:35:18 -0400
Subject: [PATCH 3/8] Add tests

---
 .../comparisonTable/ComparisonTable.test.tsx  | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
index 3c410a2b93..6719a29d96 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
@@ -400,5 +400,50 @@ describe('ComparisonTable', () => {
 
       expect(dragOverEvent.preventDefault).toHaveBeenCalled()
     })
+
+    it('should show the header being dragged in its original position until the drop', () => {
+      renderTable()
+
+      const [endingNode, startingNode] = getHeaders()
+
+      dragEnter(startingNode, endingNode.id, DragEnterDirection.LEFT)
+
+      const [, draggedHeader] = getHeaders()
+
+      expect(draggedHeader.isSameNode(startingNode)).toBe(true)
+    })
+
+    it('should wrap the drop target with the header we are dragging over', () => {
+      renderTable()
+
+      const [endingNode, startingNode] = getHeaders()
+
+      dragEnter(startingNode, endingNode.id, DragEnterDirection.LEFT)
+
+      const [headerWrapper] = getHeaders()
+
+      expect(headerWrapper.childElementCount).toBe(2)
+      expect(headerWrapper.contains(endingNode)).toBe(true)
+    })
+
+    it('should not change the order when dropping a header in its own spot', () => {
+      renderTable()
+
+      const [startingAndEndingNode, secondEndingNode] = getHeaders()
+
+      dragAndDrop(
+        startingAndEndingNode,
+        startingAndEndingNode,
+        DragEnterDirection.RIGHT
+      )
+      expect(mockPostMessage).not.toHaveBeenCalled()
+
+      dragAndDrop(
+        startingAndEndingNode,
+        secondEndingNode,
+        DragEnterDirection.RIGHT
+      )
+      expect(mockPostMessage).toHaveBeenCalled()
+    })
   })
 })

From 8315d2587dbff2e49ac860f8b72cb1fdb48ca544 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 21 Jul 2022 11:24:58 -0400
Subject: [PATCH 4/8] mend

---
 webview/src/plots/components/App.test.tsx     |  8 +--
 .../components/dragDrop/DragDropContainer.tsx | 72 +++++++++----------
 webview/src/test/dragDrop.ts                  | 21 ++----
 webview/src/test/nodes.ts                     |  3 +
 4 files changed, 46 insertions(+), 58 deletions(-)
 create mode 100644 webview/src/test/nodes.ts

diff --git a/webview/src/plots/components/App.test.tsx b/webview/src/plots/components/App.test.tsx
index 4bde293e7c..f2904bfd3a 100644
--- a/webview/src/plots/components/App.test.tsx
+++ b/webview/src/plots/components/App.test.tsx
@@ -742,7 +742,7 @@ describe('App', () => {
 
     dragEnter(
       anotherSingleViewPlot,
-      CSS.escape(movedSingleViewPlot.id),
+      movedSingleViewPlot.id,
       DragEnterDirection.LEFT
     )
 
@@ -825,7 +825,7 @@ describe('App', () => {
 
     const plots = screen.getAllByTestId(/^plot_/)
 
-    dragEnter(plots[1], CSS.escape(plots[0].id), DragEnterDirection.LEFT)
+    dragEnter(plots[1], plots[0].id, DragEnterDirection.LEFT)
 
     const plotsWithDropTarget = screen.getAllByTestId(/^plot_/)
     expect(plotsWithDropTarget.map(plot => plot.id)).toStrictEqual([
@@ -842,7 +842,7 @@ describe('App', () => {
     })
 
     const plots = screen.getAllByTestId(/^plot_/)
-    dragEnter(plots[0], CSS.escape(plots[1].id), DragEnterDirection.RIGHT)
+    dragEnter(plots[0], plots[1].id, DragEnterDirection.RIGHT)
 
     const plotsWithDropTarget = screen.getAllByTestId(/^plot_/)
 
@@ -862,7 +862,7 @@ describe('App', () => {
     const plots = screen.getAllByTestId(/^plot_/)
     expect(plots[1].style.display).not.toBe('none')
 
-    dragEnter(plots[1], CSS.escape(plots[1].id), DragEnterDirection.LEFT)
+    dragEnter(plots[1], plots[1].id, DragEnterDirection.LEFT)
 
     expect(plots[1].style.display).toBe('none')
   })
diff --git a/webview/src/shared/components/dragDrop/DragDropContainer.tsx b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
index 82d137ea4a..59c5e60ff8 100644
--- a/webview/src/shared/components/dragDrop/DragDropContainer.tsx
+++ b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
@@ -220,48 +220,44 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
     />
   )
 
+  const createItemWithDropTarget = (id: string, item: JSX.Element) => {
+    const isEnteringRight = direction === DragEnterDirection.RIGHT
+    const targetClassName = shouldShowOnDrag
+      ? cx(styles.dropTargetWhenShowingOnDrag, {
+          [styles.dropTargetWhenShowingOnDragLeft]: !isEnteringRight,
+          [styles.dropTargetWhenShowingOnDragRight]: isEnteringRight
+        })
+      : undefined
+    const target = makeTarget(
+      dropTarget,
+      handleDragOver,
+      handleOnDrop,
+      id,
+      targetClassName
+    )
+    const itemWithTag = shouldShowOnDrag ? (
+      <div key="item" {...item.props} />
+    ) : (
+      item
+    )
+    const block = isEnteringRight
+      ? [itemWithTag, target]
+      : [target, itemWithTag]
+
+    return shouldShowOnDrag ? (
+      <item.type key={item.key} className={styles.newBlockWhenShowingOnDrag}>
+        {block}
+      </item.type>
+    ) : (
+      block
+    )
+  }
+
   const wrappedItems = items.flatMap(draggable => {
     const { id } = draggable.props
     const item = id && buildItem(id, draggable)
 
-    if (id === draggedOverId) {
-      const targetClassName = shouldShowOnDrag
-        ? cx(styles.dropTargetWhenShowingOnDrag, {
-            [styles.dropTargetWhenShowingOnDragLeft]:
-              direction === DragEnterDirection.LEFT,
-            [styles.dropTargetWhenShowingOnDragRight]:
-              direction === DragEnterDirection.RIGHT
-          })
-        : undefined
-      const target = makeTarget(
-        dropTarget,
-        handleDragOver,
-        handleOnDrop,
-        id,
-        targetClassName
-      )
-      const Tag = item.type
-      const itemWithTag = shouldShowOnDrag ? (
-        <div key="item" {...item.props} />
-      ) : (
-        item
-      )
-      const block =
-        direction === DragEnterDirection.RIGHT
-          ? [itemWithTag, target]
-          : [target, itemWithTag]
-
-      if (shouldShowOnDrag) {
-        return (
-          <Tag key={draggable.key} className={styles.newBlockWhenShowingOnDrag}>
-            {block}
-          </Tag>
-        )
-      }
-      return block
-    }
-
-    return item
+    return id === draggedOverId ? createItemWithDropTarget(id, item) : item
   })
 
   const Wrapper = wrapperComponent?.component
diff --git a/webview/src/test/dragDrop.ts b/webview/src/test/dragDrop.ts
index 64750dc0d0..373a25e53b 100644
--- a/webview/src/test/dragDrop.ts
+++ b/webview/src/test/dragDrop.ts
@@ -1,4 +1,5 @@
 import { act } from 'react-dom/test-utils'
+import { idToNodeNode } from './nodes'
 import * as DragDropUtils from '../shared/components/dragDrop/util'
 
 const testStorage = new Map()
@@ -31,14 +32,12 @@ export const dragEnter = (
     jest.advanceTimersByTime(1)
   })
 
-  let draggedOver =
-    (draggedOverId && document.querySelector(`#${draggedOverId}`)) || null
+  let draggedOver = idToNodeNode(draggedOverId)
 
   act(() => {
     draggedOver?.dispatchEvent(createBubbledEvent('dragenter'))
   })
-  draggedOver =
-    (draggedOverId && document.querySelector(`#${draggedOverId}`)) || null
+  draggedOver = idToNodeNode(draggedOverId)
 
   act(() => {
     if (direction !== DragDropUtils.DragEnterDirection.AUTO) {
@@ -66,25 +65,15 @@ export const dragAndDrop = (
     .LEFT
 ) => {
   // When showing element on drag, the dragged over element is being recreacted to be wrapped in another element, thus the endingNode does not exist as is in the document
-  // CSS.escape is needed for weird ids in the experiments table (ex.: :loss)
-  const endingNodeId = CSS.escape(endingNode.id)
+  const endingNodeId = endingNode.id
   dragEnter(startingNode, endingNodeId, direction)
 
   jest.useFakeTimers()
   act(() => {
     jest.advanceTimersByTime(1)
-  })
-  act(() => {
-    const endingNodeWithId =
-      (endingNodeId && document.querySelector(`#${endingNodeId}`)) || null
+    const endingNodeWithId = idToNodeNode(endingNodeId)
     endingNodeWithId?.dispatchEvent(createBubbledEvent('drop'))
-  })
-
-  act(() => {
     jest.advanceTimersByTime(1)
-  })
-
-  act(() => {
     startingNode.dispatchEvent(createBubbledEvent('dragend'))
   })
 
diff --git a/webview/src/test/nodes.ts b/webview/src/test/nodes.ts
new file mode 100644
index 0000000000..caa58f0031
--- /dev/null
+++ b/webview/src/test/nodes.ts
@@ -0,0 +1,3 @@
+export const idToNodeNode = (id: string) =>
+  // CSS.escape is needed for weird ids (ex.: :loss, plot/1...)
+  (id && document.querySelector(`#${CSS.escape(id)}`)) || null

From 26dff253d8ea09f00717a57ebf5667c7430b6b27 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 21 Jul 2022 11:42:04 -0400
Subject: [PATCH 5/8] Remove unused style

---
 webview/src/shared/components/dragDrop/styles.module.scss | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/webview/src/shared/components/dragDrop/styles.module.scss b/webview/src/shared/components/dragDrop/styles.module.scss
index 878b7e6d5d..1e76541e71 100644
--- a/webview/src/shared/components/dragDrop/styles.module.scss
+++ b/webview/src/shared/components/dragDrop/styles.module.scss
@@ -16,7 +16,3 @@
 .dropTargetWhenShowingOnDragRight {
   right: -2px;
 }
-
-.hiddenOnDrag {
-  display: none;
-}

From 16bb5259c71ceea877c556898a0956314b1722db Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 21 Jul 2022 12:00:11 -0400
Subject: [PATCH 6/8] Try to fix windows tests

---
 webview/src/test/nodes.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/webview/src/test/nodes.ts b/webview/src/test/nodes.ts
index caa58f0031..87032ace92 100644
--- a/webview/src/test/nodes.ts
+++ b/webview/src/test/nodes.ts
@@ -1,3 +1,3 @@
 export const idToNodeNode = (id: string) =>
-  // CSS.escape is needed for weird ids (ex.: :loss, plot/1...)
-  (id && document.querySelector(`#${CSS.escape(id)}`)) || null
+  // eslint-disable-next-line unicorn/prefer-query-selector
+  (id && document.getElementById(id)) || null

From 5ccda07b74ebd52ea351892a4dbcdfd70c58a84f Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 21 Jul 2022 14:34:40 -0400
Subject: [PATCH 7/8] Change the style of the dragged column in the comparison
 table

---
 .../components/comparisonTable/ComparisonTableHead.tsx   | 9 ++++++++-
 .../components/comparisonTable/ComparisonTableRow.tsx    | 8 +++++++-
 .../plots/components/comparisonTable/styles.module.scss  | 4 ++++
 .../src/shared/components/dragDrop/DragDropContainer.tsx | 8 +++++++-
 4 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx
index aeb9c09513..4942215865 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableHead.tsx
@@ -1,10 +1,12 @@
 import React from 'react'
 import { Revision } from 'dvc/src/plots/webview/contract'
 import cx from 'classnames'
+import { useSelector } from 'react-redux'
 import styles from './styles.module.scss'
 import { DropTarget } from './DropTarget'
 import { ComparisonTableHeader } from './ComparisonTableHeader'
 import { DragDropContainer } from '../../../shared/components/dragDrop/DragDropContainer'
+import { RootState } from '../../store'
 
 export type ComparisonTableColumn = Revision
 
@@ -21,6 +23,10 @@ export const ComparisonTableHead: React.FC<ComparisonTableHeadProps> = ({
   setColumnsOrder,
   setPinnedColumn
 }) => {
+  const draggedId = useSelector(
+    (state: RootState) => state.dragAndDrop.draggedRef?.itemId
+  )
+
   const items = columns.map(({ revision, displayColor, group }) => {
     const isPinned = revision === pinnedColumn
     return (
@@ -28,7 +34,8 @@ export const ComparisonTableHead: React.FC<ComparisonTableHeadProps> = ({
         key={revision}
         id={revision}
         className={cx(styles.comparisonTableHeader, {
-          [styles.pinnedColumnHeader]: isPinned
+          [styles.pinnedColumnHeader]: isPinned,
+          [styles.draggedColumn]: draggedId === revision
         })}
       >
         <ComparisonTableHeader
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx
index 9527df50cd..1bfe353dbf 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx
@@ -2,11 +2,13 @@ import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 import React, { useState } from 'react'
 import cx from 'classnames'
+import { useSelector } from 'react-redux'
 import styles from './styles.module.scss'
 import { Icon } from '../../../shared/components/Icon'
 import { RefreshButton } from '../../../shared/components/button/RefreshButton'
 import { sendMessage } from '../../../shared/vscode'
 import { ChevronDown, ChevronRight } from '../../../shared/components/icons'
+import { RootState } from '../../store'
 
 export interface ComparisonTableRowProps {
   path: string
@@ -21,6 +23,9 @@ export const ComparisonTableRow: React.FC<ComparisonTableRowProps> = ({
   nbColumns,
   pinnedColumn
 }) => {
+  const draggedId = useSelector(
+    (state: RootState) => state.dragAndDrop.draggedRef?.itemId
+  )
   const [isShown, setIsShown] = useState(true)
 
   const toggleIsShownState = () => setIsShown(!isShown)
@@ -46,7 +51,8 @@ export const ComparisonTableRow: React.FC<ComparisonTableRowProps> = ({
               key={path + plot.revision}
               className={cx({
                 [styles.pinnedColumnCell]: isPinned,
-                [styles.missing]: isShown && missing
+                [styles.missing]: isShown && missing,
+                [styles.draggedColumn]: draggedId === plot.revision
               })}
             >
               <div
diff --git a/webview/src/plots/components/comparisonTable/styles.module.scss b/webview/src/plots/components/comparisonTable/styles.module.scss
index 95235ef63e..d858f2dba1 100644
--- a/webview/src/plots/components/comparisonTable/styles.module.scss
+++ b/webview/src/plots/components/comparisonTable/styles.module.scss
@@ -163,3 +163,7 @@ $gap: 4px;
   height: 100%;
   border-right: 2px dashed $accent-color;
 }
+
+.draggedColumn {
+  opacity: 0.4;
+}
diff --git a/webview/src/shared/components/dragDrop/DragDropContainer.tsx b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
index 59c5e60ff8..f07e94239c 100644
--- a/webview/src/shared/components/dragDrop/DragDropContainer.tsx
+++ b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
@@ -157,6 +157,7 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
       setDraggedId('')
     }, 0)
     if (e.dataTransfer.getData('itemId') === draggedOverId) {
+      dispatch(changeRef(undefined))
       return
     }
     const newOrder = [...order]
@@ -203,12 +204,17 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
     }
   }
 
+  const handleDragEnd = () => {
+    dispatch(changeRef(undefined))
+    cleanup()
+  }
+
   const buildItem = (id: string, draggable: JSX.Element) => (
     <draggable.type
       key={draggable.key}
       {...draggable.props}
       onDragStart={handleDragStart}
-      onDragEnd={cleanup}
+      onDragEnd={handleDragEnd}
       onDragOver={handleDragOver}
       onDragEnter={handleDragEnter}
       onDrop={handleOnDrop}

From cb9df126f05146047250e032bd653749cad12c45 Mon Sep 17 00:00:00 2001
From: Stephanie Roy <stephanie@iterative.ai>
Date: Thu, 21 Jul 2022 14:38:19 -0400
Subject: [PATCH 8/8] Fix test

---
 .../comparisonTable/ComparisonTableRow.test.tsx   | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRow.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRow.test.tsx
index b6dd47a138..b27543f77a 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRow.test.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRow.test.tsx
@@ -1,16 +1,19 @@
 /**
  * @jest-environment jsdom
  */
+import { configureStore } from '@reduxjs/toolkit'
 import { join } from 'dvc/src/test/util/path'
 import React from 'react'
 import { render, cleanup, screen, fireEvent } from '@testing-library/react'
 import '@testing-library/jest-dom/extend-expect'
 import comparisonPlotsFixture from 'dvc/src/test/fixtures/plotsDiff/comparison'
+import { Provider } from 'react-redux'
 import {
   ComparisonTableRow,
   ComparisonTableRowProps
 } from './ComparisonTableRow'
 import styles from '../styles.module.scss'
+import { storeReducers } from '../../store'
 
 jest.mock('../../../shared/api')
 
@@ -32,9 +35,15 @@ describe('ComparisonTableRow', () => {
 
   const renderRow = (props = basicProps) =>
     render(
-      <table>
-        <ComparisonTableRow {...props} />
-      </table>
+      <Provider
+        store={configureStore({
+          reducer: storeReducers
+        })}
+      >
+        <table>
+          <ComparisonTableRow {...props} />
+        </table>
+      </Provider>
     )
 
   it('should render a row toggler', () => {