From 83b831784dea38677775482953b89987cdd6969d Mon Sep 17 00:00:00 2001
From: Illia Olenchenko <icrosil@gmail.com>
Date: Thu, 16 Nov 2023 14:51:06 +0000
Subject: [PATCH] Fixed column header on merged cells (#5230)

---
 .../__tests__/e2e/Tables.spec.mjs             | 95 +++++++++++++++++++
 .../__tests__/utils/index.mjs                 |  5 +
 .../plugins/TableActionMenuPlugin/index.tsx   | 18 +++-
 3 files changed, 114 insertions(+), 4 deletions(-)

diff --git a/packages/lexical-playground/__tests__/e2e/Tables.spec.mjs b/packages/lexical-playground/__tests__/e2e/Tables.spec.mjs
index 8551f686d19..db94b26e261 100644
--- a/packages/lexical-playground/__tests__/e2e/Tables.spec.mjs
+++ b/packages/lexical-playground/__tests__/e2e/Tables.spec.mjs
@@ -38,6 +38,7 @@ import {
   selectFromAdditionalStylesDropdown,
   setBackgroundColor,
   test,
+  toggleColumnHeader,
   unmergeTableCell,
 } from '../utils/index.mjs';
 
@@ -1961,4 +1962,98 @@ test.describe('Tables', () => {
       `,
     );
   });
+
+  test('Add column header after merging cells #4378', async ({
+    page,
+    isPlainText,
+    isCollab,
+  }) => {
+    await initialize({isCollab, page});
+    test.skip(isPlainText);
+    if (IS_COLLAB) {
+      // The contextual menu positioning needs fixing (it's hardcoded to show on the right side)
+      page.setViewportSize({height: 1000, width: 3000});
+    }
+
+    await focusEditor(page);
+
+    await insertTable(page, 4, 4);
+    await selectCellsFromTableCords(
+      page,
+      {x: 1, y: 2},
+      {x: 3, y: 3},
+      false,
+      false,
+    );
+    await mergeTableCells(page);
+    await selectCellsFromTableCords(
+      page,
+      {x: 3, y: 1},
+      {x: 3, y: 1},
+      false,
+      false,
+    );
+    await toggleColumnHeader(page);
+    await assertHTML(
+      page,
+      html`
+        <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+        <table class="PlaygroundEditorTheme__table">
+          <tr>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+          </tr>
+          <tr>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+            <td class="PlaygroundEditorTheme__tableCell">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </td>
+            <td class="PlaygroundEditorTheme__tableCell">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </td>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+          </tr>
+          <tr>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+            <td
+              class="PlaygroundEditorTheme__tableCell"
+              colspan="3"
+              rowspan="2">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </td>
+          </tr>
+          <tr>
+            <th
+              class="PlaygroundEditorTheme__tableCell PlaygroundEditorTheme__tableCellHeader">
+              <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+            </th>
+          </tr>
+        </table>
+        <p class="PlaygroundEditorTheme__paragraph"><br /></p>
+      `,
+    );
+  });
 });
diff --git a/packages/lexical-playground/__tests__/utils/index.mjs b/packages/lexical-playground/__tests__/utils/index.mjs
index 9b5d4567d59..0d199f21ee7 100644
--- a/packages/lexical-playground/__tests__/utils/index.mjs
+++ b/packages/lexical-playground/__tests__/utils/index.mjs
@@ -881,6 +881,11 @@ export async function unmergeTableCell(page) {
   await click(page, '.item[data-test-id="table-unmerge-cells"]');
 }
 
+export async function toggleColumnHeader(page) {
+  await click(page, '.table-cell-action-button-container');
+  await click(page, '.item[data-test-id="table-column-header"]');
+}
+
 export async function deleteTableRows(page) {
   await click(page, '.table-cell-action-button-container');
   await click(page, '.item[data-test-id="table-delete-rows"]');
diff --git a/packages/lexical-playground/src/plugins/TableActionMenuPlugin/index.tsx b/packages/lexical-playground/src/plugins/TableActionMenuPlugin/index.tsx
index 267366fde13..48f47d8667c 100644
--- a/packages/lexical-playground/src/plugins/TableActionMenuPlugin/index.tsx
+++ b/packages/lexical-playground/src/plugins/TableActionMenuPlugin/index.tsx
@@ -432,6 +432,13 @@ function TableActionMenu({
         $getTableColumnIndexFromTableCellNode(tableCellNode);
 
       const tableRows = tableNode.getChildren();
+      const maxRowsLength = Math.max(
+        ...tableRows.map((row) => row.getChildren().length),
+      );
+
+      if (tableColumnIndex >= maxRowsLength || tableColumnIndex < 0) {
+        throw new Error('Expected table cell to be inside of table row.');
+      }
 
       for (let r = 0; r < tableRows.length; r++) {
         const tableRow = tableRows[r];
@@ -441,9 +448,9 @@ function TableActionMenu({
         }
 
         const tableCells = tableRow.getChildren();
-
-        if (tableColumnIndex >= tableCells.length || tableColumnIndex < 0) {
-          throw new Error('Expected table cell to be inside of table row.');
+        if (tableColumnIndex >= tableCells.length) {
+          // if cell is outside of bounds for the current row (for example various merge cell cases) we shouldn't highlight it
+          continue;
         }
 
         const tableCell = tableCells[tableColumnIndex];
@@ -598,7 +605,10 @@ function TableActionMenu({
           row header
         </span>
       </button>
-      <button className="item" onClick={() => toggleTableColumnIsHeader()}>
+      <button
+        className="item"
+        onClick={() => toggleTableColumnIsHeader()}
+        data-test-id="table-column-header">
         <span className="text">
           {(tableCellNode.__headerState & TableCellHeaderStates.COLUMN) ===
           TableCellHeaderStates.COLUMN