From bb0132cca3abf90f275b213859fed83a5fa06fb2 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Fri, 1 Dec 2023 11:38:04 +1100
Subject: [PATCH] Handle missing items in drag and drop container

---
 .../comparisonTable/ComparisonTable.test.tsx  | 19 +++++++
 .../comparisonTable/ComparisonTableRows.tsx   | 54 ++++++++++---------
 .../components/dragDrop/DragDropContainer.tsx |  3 ++
 3 files changed, 50 insertions(+), 26 deletions(-)

diff --git a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
index d9674d7b5c..52d7d3613c 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTable.test.tsx
@@ -388,6 +388,25 @@ describe('ComparisonTable', () => {
     expect(aHeader.style.top).not.toBe('')
   })
 
+  it('should not throw an error when an image is removed from the data', () => {
+    const { rerender } = renderTable()
+
+    const plotsWithMissingImage = comparisonTableFixture.plots.slice(1)
+    expect(plotsWithMissingImage.length).toStrictEqual(
+      comparisonTableFixture.plots.length - 1
+    )
+
+    expect(() =>
+      renderTable(
+        {
+          ...comparisonTableFixture,
+          plots: plotsWithMissingImage
+        },
+        rerender
+      )
+    ).not.toThrow()
+  })
+
   describe('Columns drag and drop', () => {
     const pinSecondColumn = () => {
       const secondColumn = getPin(screen.getByText(revisions[1]))
diff --git a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
index 045290e6e6..ab861809a2 100644
--- a/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
+++ b/webview/src/plots/components/comparisonTable/ComparisonTableRows.tsx
@@ -37,31 +37,33 @@ export const ComparisonTableRows: React.FC<ComparisonTableRowsProps> = ({
     dispatch(changeRowHeight(firstRowHeight))
   }
 
-  const rows = rowsOrder.map((path, i) => {
-    const plot = plots.find(p => p.path === path)
-    if (!plot) {
-      return
-    }
-    const revs = plot.revisions
-    return (
-      <tbody
-        data-testid="comparison-table-body"
-        key={path}
-        id={path}
-        ref={i === 0 ? firstRowRef : undefined}
-      >
-        <ComparisonTableRow
-          path={path}
-          plots={columns.map(column => ({
-            id: column.id,
-            imgs: revs[column.id]?.imgs
-          }))}
-          nbColumns={columns.length}
-          pinnedColumn={pinnedColumn}
-        />
-      </tbody>
-    )
-  })
+  const rows = rowsOrder
+    .map((path, i) => {
+      const plot = plots.find(p => p.path === path)
+      if (!plot) {
+        return
+      }
+      const revs = plot.revisions
+      return (
+        <tbody
+          data-testid="comparison-table-body"
+          key={path}
+          id={path}
+          ref={i === 0 ? firstRowRef : undefined}
+        >
+          <ComparisonTableRow
+            path={path}
+            plots={columns.map(column => ({
+              id: column.id,
+              imgs: revs[column.id]?.imgs
+            }))}
+            nbColumns={columns.length}
+            pinnedColumn={pinnedColumn}
+          />
+        </tbody>
+      )
+    })
+    .filter(Boolean) as JSX.Element[]
 
   const changeRowsOrder = (order: string[]) => {
     setRowsOrder(order)
@@ -70,7 +72,7 @@ export const ComparisonTableRows: React.FC<ComparisonTableRowsProps> = ({
 
   return (
     <DragDropContainer
-      items={rows as JSX.Element[]}
+      items={rows}
       order={rowsOrder}
       setOrder={changeRowsOrder}
       group="comparison-table"
diff --git a/webview/src/shared/components/dragDrop/DragDropContainer.tsx b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
index 401e9b11e4..9a1b54bbaa 100644
--- a/webview/src/shared/components/dragDrop/DragDropContainer.tsx
+++ b/webview/src/shared/components/dragDrop/DragDropContainer.tsx
@@ -297,6 +297,9 @@ export const DragDropContainer: React.FC<DragDropContainerProps> = ({
 
   const wrappedItems = items
     .map(draggable => {
+      if (!draggable) {
+        return
+      }
       const id = draggable.props.id
       const isDraggedOver =
         id === draggedOverId && (hoveringSomething || !parentDraggedOver)