[lexical-table][lexical-y] Bug Fix: Don't throw when a table exists before its DOM

.changeset/table-observer-unrendered-table.md

@@ -0,0 +1,5 @@
+---
+'@lexical/table': patch
+---
+
+[fix] Building an editor no longer throws `Expected to find TableElement in DOM` when the editor state already holds a table before `TableExtension` registers, for example when an extension listed earlier (such as `YCollabExtension` on a document that holds a table) sets the content in its `register`; the table's selection handling attaches when the table is rendered

packages/lexical-table/src/LexicalTablePluginHelpers.ts

@@ -362,6 +362,13 @@ export function registerTableSelectionObserver(
           for (const [nodeKey, mutation] of nodeMutations) {
             const tableSelection = tableObservers.observers.get(nodeKey);
             if (mutation === 'created' || mutation === 'updated') {
+              if (editor.getElementByKey(nodeKey) === null) {
+                // The table is not rendered, as when the editor has no root
+                // element yet. The reconciler reports it as created when it
+                // renders it, and the observer attaches then.
+                tableObservers.removeObserver(nodeKey);
+                continue;
+              }
               const {tableNode, tableElement} =
                 $getTableAndElementByKey(nodeKey);
               if (tableSelection === undefined) {

packages/lexical-table/src/__tests__/unit/TableSelectionObserverBeforeDOM.test.ts

@@ -0,0 +1,149 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ */
+import {buildEditorFromExtensions} from '@lexical/extension';
+import {
+  $createTableNodeWithDimensions,
+  $isTableNode,
+  getTableObserverFromTableElement,
+  type HTMLTableElementWithWithTableSelectionState,
+  registerTablePlugin,
+  registerTableSelectionObserver,
+  TableCellNode,
+  TableExtension,
+  TableNode,
+  TableRowNode,
+} from '@lexical/table';
+import {
+  $createParagraphNode,
+  $getRoot,
+  createEditor,
+  defineExtension,
+  type LexicalEditor,
+  type LexicalEditorWithDispose,
+  mergeRegister,
+} from 'lexical';
+import {afterEach, assert, describe, expect, test} from 'vitest';
+
+function $appendTable() {
+  $getRoot().append(
+    $createTableNodeWithDimensions(2, 2, false),
+    $createParagraphNode(),
+  );
+}
+
+function $getTableKey(): string | null {
+  const table = $getRoot().getChildren().find($isTableNode);
+  return table ? table.getKey() : null;
+}
+
+/** Serialized editor state holding one 2x2 table, made by a separate editor. */
+function tableStateJSON(): string {
+  const source = buildEditorFromExtensions(
+    defineExtension({dependencies: [TableExtension], name: 'table-source'}),
+  );
+  source.update($appendTable, {discrete: true});
+  const json = JSON.stringify(source.getEditorState().toJSON());
+  source.dispose();
+  return json;
+}
+
+function getRenderedTable(
+  rootElement: HTMLElement,
+): HTMLTableElementWithWithTableSelectionState {
+  const tableElement = rootElement.querySelector('table');
+  assert(tableElement !== null, 'Expected a rendered table');
+  return tableElement;
+}
+
+/**
+ * An extension listed before TableExtension that replaces the editor state
+ * from its register, as an extension that loads stored content does.
+ */
+function setStateOnRegister(json: string) {
+  return defineExtension({
+    name: 'set-state-on-register',
+    register: (editor: LexicalEditor) => {
+      editor.setEditorState(editor.parseEditorState(json));
+      return () => {};
+    },
+  });
+}
+
+describe('table selection observer for tables that exist before their DOM', () => {
+  const cleanups: (() => void)[] = [];
+  afterEach(() => {
+    for (const cleanup of cleanups.splice(0).reverse()) {
+      cleanup();
+    }
+  });
+
+  function buildWithTableSetOnRegister(): LexicalEditorWithDispose {
+    const editor = buildEditorFromExtensions(
+      defineExtension({
+        dependencies: [setStateOnRegister(tableStateJSON()), TableExtension],
+        name: 'state-before-table',
+      }),
+    );
+    cleanups.push(() => editor.dispose());
+    return editor;
+  }
+
+  function attachRoot(editor: LexicalEditor): HTMLDivElement {
+    const rootElement = document.createElement('div');
+    rootElement.contentEditable = 'true';
+    document.body.appendChild(rootElement);
+    cleanups.push(() => rootElement.remove());
+    editor.setRootElement(rootElement);
+    cleanups.push(() => editor.setRootElement(null));
+    return rootElement;
+  }
+
+  test('builds when an earlier extension sets a state holding a table during register', () => {
+    const editor = buildWithTableSetOnRegister();
+    expect(editor.read($getTableKey)).not.toBe(null);
+  });
+
+  test('observes the table once a root element renders it', () => {
+    const editor = buildWithTableSetOnRegister();
+    const rootElement = attachRoot(editor);
+    const tableElement = getRenderedTable(rootElement);
+    const observer = getTableObserverFromTableElement(tableElement);
+    assert(observer !== null, 'Expected the rendered table to be observed');
+    expect(observer.tableNodeKey).toBe(editor.read($getTableKey));
+  });
+
+  test('observes the new table DOM when the root element is replaced', () => {
+    const editor = buildWithTableSetOnRegister();
+    const firstTable = getRenderedTable(attachRoot(editor));
+    editor.setRootElement(null);
+    const secondTable = getRenderedTable(attachRoot(editor));
+    expect(secondTable).not.toBe(firstTable);
+    expect(getTableObserverFromTableElement(firstTable)).toBe(null);
+    expect(getTableObserverFromTableElement(secondTable)).not.toBe(null);
+  });
+
+  test('registerTableSelectionObserver accepts an unrendered table', () => {
+    const editor = createEditor({
+      nodes: [TableNode, TableRowNode, TableCellNode],
+      onError: error => {
+        throw error;
+      },
+    });
+    editor.update($appendTable, {discrete: true});
+    cleanups.push(
+      mergeRegister(
+        registerTablePlugin(editor),
+        registerTableSelectionObserver(editor),
+      ),
+    );
+    const rootElement = attachRoot(editor);
+    expect(
+      getTableObserverFromTableElement(getRenderedTable(rootElement)),
+    ).not.toBe(null);
+  });
+});

packages/lexical-y/src/__tests__/unit/TableOnOpen.test.ts

@@ -0,0 +1,75 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ */
+
+// @vitest-environment node
+
+// An editor bound to a document that already holds a table opens with the
+// table, wherever YCollabExtension and TableExtension sit in its
+// dependencies. The binding reads the document when it registers, so with
+// YCollabExtension listed first the table is in the editor state before
+// TableExtension registers, and before there is a root element.
+
+import {buildEditorFromExtensions} from '@lexical/extension';
+import {RichTextExtension} from '@lexical/rich-text';
+import {
+  $createTableCellNode,
+  $createTableNode,
+  $createTableRowNode,
+  TableExtension,
+} from '@lexical/table';
+import * as Y from '@y/y';
+import {
+  $createParagraphNode,
+  $createTextNode,
+  $getRoot,
+  configExtension,
+  defineExtension,
+} from 'lexical';
+import {expect, test} from 'vitest';
+
+import {YCollabExtension} from '../..';
+import {seedYState} from '../../headless';
+
+function $appendTable() {
+  $getRoot().append(
+    $createTableNode().append(
+      $createTableRowNode().append(
+        $createTableCellNode().append(
+          $createParagraphNode().append($createTextNode('cell')),
+        ),
+      ),
+    ),
+    $createParagraphNode(),
+  );
+}
+
+test.each(['YCollabExtension', 'TableExtension'] as const)(
+  'opens a document that holds a table with %s listed first',
+  first => {
+    const doc = new Y.Doc();
+    Y.applyUpdate(
+      doc,
+      seedYState(
+        {extensions: [RichTextExtension, TableExtension]},
+        $appendTable,
+      ),
+    );
+    const collab = configExtension(YCollabExtension, {doc});
+    const editor = buildEditorFromExtensions(
+      defineExtension({
+        dependencies:
+          first === 'YCollabExtension'
+            ? [RichTextExtension, collab, TableExtension]
+            : [RichTextExtension, TableExtension, collab],
+        name: '@lexical/y/test/table-on-open',
+      }),
+    );
+    expect(editor.read(() => $getRoot().getTextContent())).toContain('cell');
+    editor.dispose();
+  },
+);