reactive context menu overrides (#2697)

Previously, we were calling context menu `overrides` in a `useMemo`, so
they weren't updating reactively in the way that most of our other
schema overrides do. This diff calls `override` in a `useValue` instead
so it updates reactively.

It also fixes some issues with testing the `<Tldraw />` component:
currently we get a lot of errors in the console about updates not being
wrapped in `act`. These are caused by the fill patterns at different
zoom levels popping in without us waiting for them. Now, we have a
helper for rendering the tldraw component that waits for this correctly
and stops the error.

### Change Type

- [x] `patch` — Bug fix

### Test Plan


- [x] Unit Tests


### Release Notes

- Context Menu overrides will now update reactively

packages/tldraw/setupTests.js

@@ -55,3 +55,15 @@ window.fetch = async (input, init) => {
 
 	throw new Error(`Unhandled request: ${input}`)
 }
+
+window.DOMRect = class DOMRect {
+	static fromRect(rect) {
+		return new DOMRect(rect.x, rect.y, rect.width, rect.height)
+	}
+	constructor(x, y, width, height) {
+		this.x = x
+		this.y = y
+		this.width = width
+		this.height = height
+	}
+}

packages/tldraw/src/lib/Tldraw.test.tsx

@@ -1,16 +1,15 @@
-import { act, render, screen } from '@testing-library/react'
+import { act, screen } from '@testing-library/react'
 import { BaseBoxShapeUtil, Editor } from '@tldraw/editor'
 import { useState } from 'react'
+import { renderTldrawComponent } from '../test/testutils/renderTldrawComponent'
 import { Tldraw } from './Tldraw'
 
 describe('<Tldraw />', () => {
 	it('Renders without crashing', async () => {
-		await act(async () =>
-			render(
-				<Tldraw>
-					<div data-testid="canvas-1" />
-				</Tldraw>
-			)
+		await renderTldrawComponent(
+			<Tldraw>
+				<div data-testid="canvas-1" />
+			</Tldraw>
 		)
 
 		await screen.findByTestId('canvas-1')
@@ -27,7 +26,7 @@ describe('<Tldraw />', () => {
 			)
 		}
 
-		await act(async () => render(<TestComponent />))
+		await renderTldrawComponent(<TestComponent />)
 		await screen.findByTestId('canvas-1')
 	})
 
@@ -57,13 +56,12 @@ describe('<Tldraw />', () => {
 			}
 		}
 
-		const rendered = await act(async () =>
-			render(
-				<Tldraw shapeUtils={[FakeShapeUtil1]}>
-					<div data-testid="canvas-1" />
-				</Tldraw>
-			)
+		const rendered = await renderTldrawComponent(
+			<Tldraw shapeUtils={[FakeShapeUtil1]}>
+				<div data-testid="canvas-1" />
+			</Tldraw>
 		)
+
 		await screen.findByTestId('canvas-1')
 
 		await act(async () => {

packages/tldraw/src/lib/shapes/shared/defaultStyleDefs.tsx

@@ -266,7 +266,11 @@ function PatternFillDefForCanvas() {
 		}
 	}, [editor, isReady])
 
-	return <g ref={containerRef}>{defs}</g>
+	return (
+		<g ref={containerRef} data-testid={isReady ? 'ready-pattern-fill-defs' : undefined}>
+			{defs}
+		</g>
+	)
 }
 
 function findHtmlLayerParent(element: Element): HTMLElement | null {

packages/tldraw/src/lib/ui/components/ContextMenu.tsx

@@ -221,6 +221,7 @@ const ContextMenuContent = forwardRef(function ContextMenuContent() {
 		<_ContextMenu.Portal container={container}>
 			<_ContextMenu.Content
 				className="tlui-menu scrollable"
+				data-testid="context-menu"
 				alignOffset={-4}
 				collisionPadding={4}
 				onContextMenu={preventDefault}

packages/tldraw/src/lib/ui/hooks/useContextMenuSchema.tsx

@@ -87,8 +87,8 @@ export const TLUiContextMenuSchemaProvider = track(function TLUiContextMenuSchem
 		editor.getSortedChildIdsForParent(onlySelectedShape).length > 0
 	const isShapeLocked = onlySelectedShape && editor.isShapeOrAncestorLocked(onlySelectedShape)
 
-	const contextTLUiMenuSchema = useMemo<TLUiMenuSchema>(() => {
-		let contextTLUiMenuSchema: TLUiContextTTLUiMenuSchemaContextType = compactMenuItems([
+	const contextTLUiMenuSchemaWithoutOverrides = useMemo<TLUiMenuSchema>(() => {
+		return compactMenuItems([
 			menuGroup(
 				'selection',
 				showAutoSizeToggle && menuItem(actions['toggle-auto-size']),
@@ -204,23 +204,7 @@ export const TLUiContextMenuSchemaProvider = track(function TLUiContextMenuSchem
 				),
 			oneSelected && !isShapeLocked && menuGroup('delete-group', menuItem(actions['delete'])),
 		])
-
-		if (overrides) {
-			contextTLUiMenuSchema = overrides(editor, contextTLUiMenuSchema, {
-				actions,
-				oneSelected,
-				twoSelected,
-				threeSelected,
-				showAutoSizeToggle,
-				showUngroup: allowUngroup,
-				onlyFlippableShapeSelected,
-			})
-		}
-
-		return contextTLUiMenuSchema
 	}, [
-		editor,
-		overrides,
 		actions,
 		oneSelected,
 		twoSelected,
@@ -241,6 +225,34 @@ export const TLUiContextMenuSchemaProvider = track(function TLUiContextMenuSchem
 		isShapeLocked,
 	])
 
+	const contextTLUiMenuSchema = useValue(
+		'overrides',
+		() => {
+			if (!overrides) return contextTLUiMenuSchemaWithoutOverrides
+			return overrides(editor, contextTLUiMenuSchemaWithoutOverrides, {
+				actions,
+				oneSelected,
+				twoSelected,
+				threeSelected,
+				showAutoSizeToggle,
+				showUngroup: allowUngroup,
+				onlyFlippableShapeSelected,
+			})
+		},
+		[
+			actions,
+			allowUngroup,
+			contextTLUiMenuSchemaWithoutOverrides,
+			editor,
+			oneSelected,
+			onlyFlippableShapeSelected,
+			overrides,
+			showAutoSizeToggle,
+			threeSelected,
+			twoSelected,
+		]
+	)
+
 	return (
 		<TLUiContextMenuSchemaContext.Provider value={contextTLUiMenuSchema}>
 			{children}

packages/tldraw/src/test/TldrawEditor.test.tsx

@@ -13,6 +13,7 @@ import {
 } from '@tldraw/editor'
 import { defaultTools } from '../lib/defaultTools'
 import { GeoShapeUtil } from '../lib/shapes/geo/GeoShapeUtil'
+import { renderTldrawComponent } from './testutils/renderTldrawComponent'
 
 function checkAllShapes(editor: Editor, shapes: string[]) {
 	expect(Object.keys(editor!.store.schema.types.shape.migrations.subTypeMigrations!)).toStrictEqual(
@@ -24,17 +25,19 @@ function checkAllShapes(editor: Editor, shapes: string[]) {
 
 describe('<TldrawEditor />', () => {
 	it('Renders without crashing', async () => {
-		render(
+		await renderTldrawComponent(
 			<TldrawEditor tools={defaultTools} autoFocus initialState="select">
 				<div data-testid="canvas-1" />
-			</TldrawEditor>
+				<Canvas />
+			</TldrawEditor>,
+			{ waitForPatterns: false }
 		)
 		await screen.findByTestId('canvas-1')
 	})
 
 	it('Creates its own store with core shapes', async () => {
 		let editor: Editor
-		render(
+		await renderTldrawComponent(
 			<TldrawEditor
 				onMount={(e) => {
 					editor = e
@@ -44,7 +47,8 @@ describe('<TldrawEditor />', () => {
 				autoFocus
 			>
 				<div data-testid="canvas-1" />
-			</TldrawEditor>
+			</TldrawEditor>,
+			{ waitForPatterns: false }
 		)
 		await screen.findByTestId('canvas-1')
 		checkAllShapes(editor!, ['group'])
@@ -52,7 +56,7 @@ describe('<TldrawEditor />', () => {
 
 	it('Can be created with default shapes', async () => {
 		let editor: Editor
-		render(
+		await renderTldrawComponent(
 			<TldrawEditor
 				shapeUtils={[]}
 				tools={defaultTools}
@@ -63,7 +67,9 @@ describe('<TldrawEditor />', () => {
 				autoFocus
 			>
 				<div data-testid="canvas-1" />
-			</TldrawEditor>
+				<Canvas />
+			</TldrawEditor>,
+			{ waitForPatterns: false }
 		)
 		await screen.findByTestId('canvas-1')
 		expect(editor!).toBeTruthy()
@@ -73,7 +79,7 @@ describe('<TldrawEditor />', () => {
 
 	it('Renders with an external store', async () => {
 		const store = createTLStore({ shapeUtils: [] })
-		render(
+		await renderTldrawComponent(
 			<TldrawEditor
 				store={store}
 				tools={defaultTools}
@@ -84,7 +90,9 @@ describe('<TldrawEditor />', () => {
 				autoFocus
 			>
 				<div data-testid="canvas-1" />
-			</TldrawEditor>
+				<Canvas />
+			</TldrawEditor>,
+			{ waitForPatterns: false }
 		)
 		await screen.findByTestId('canvas-1')
 	})
@@ -184,21 +192,19 @@ describe('<TldrawEditor />', () => {
 
 	it('Renders the canvas and shapes', async () => {
 		let editor = {} as Editor
-		await act(async () =>
-			render(
-				<TldrawEditor
-					shapeUtils={[GeoShapeUtil]}
-					initialState="select"
-					tools={defaultTools}
-					autoFocus
-					onMount={(editorApp) => {
-						editor = editorApp
-					}}
-				>
-					<Canvas />
-					<div data-testid="canvas-1" />
-				</TldrawEditor>
-			)
+		await renderTldrawComponent(
+			<TldrawEditor
+				shapeUtils={[GeoShapeUtil]}
+				initialState="select"
+				tools={defaultTools}
+				autoFocus
+				onMount={(editorApp) => {
+					editor = editorApp
+				}}
+			>
+				<Canvas />
+				<div data-testid="canvas-1" />
+			</TldrawEditor>
 		)
 		await screen.findByTestId('canvas-1')
 
@@ -310,7 +316,7 @@ describe('Custom shapes', () => {
 
 	it('Uses custom shapes', async () => {
 		let editor = {} as Editor
-		render(
+		await renderTldrawComponent(
 			<TldrawEditor
 				shapeUtils={shapeUtils}
 				tools={[...defaultTools, ...tools]}
@@ -322,7 +328,8 @@ describe('Custom shapes', () => {
 			>
 				<Canvas />
 				<div data-testid="canvas-1" />
-			</TldrawEditor>
+			</TldrawEditor>,
+			{ waitForPatterns: false }
 		)
 		await screen.findByTestId('canvas-1')
 

packages/tldraw/src/test/testutils/renderTldrawComponent.tsx

@@ -0,0 +1,36 @@
+import { render } from '@testing-library/react'
+import { Editor, promiseWithResolve } from '@tldraw/editor'
+import { ReactElement } from 'react'
+
+/**
+ * Render a <Tldraw /> component and wait for it to become ready. By default, this includes waiting
+ * for the canvas and fill patterns to pre-render.
+ *
+ * You almost always want to wait for the canvas, but if you are testing a component that passes its
+ * own children and doesn't render the canvas itself, you can set `waitForCanvas` to false.
+ *
+ * Without waiting for patterns, a bunch of "missing `act()`" errors will fill the console, but if
+ * you don't need it (or your're testing the tldraw component without our default shapes, and so
+ * don't have pre-rendered patterns to worry about) you can set `waitForPatterns` to false.
+ */
+export async function renderTldrawComponent(
+	element: ReactElement,
+	{ waitForPatterns = true } = {}
+) {
+	const result = render(element)
+	if (waitForPatterns) await result.findByTestId('ready-pattern-fill-defs')
+	return result
+}
+
+export async function renderTldrawComponentWithEditor(
+	cb: (onMount: (editor: Editor) => void) => ReactElement,
+	opts?: { waitForPatterns?: boolean }
+) {
+	const editorPromise = promiseWithResolve<Editor>()
+	const element = cb((editor) => {
+		editorPromise.resolve(editor)
+	})
+	const rendered = await renderTldrawComponent(element, opts)
+	const editor = await editorPromise
+	return { editor, rendered }
+}

packages/tldraw/src/test/ui/ContextMenu.test.tsx

@@ -0,0 +1,74 @@
+import { fireEvent, screen } from '@testing-library/react'
+import { createShapeId, noop } from '@tldraw/editor'
+import { act } from 'react-dom/test-utils'
+import { Tldraw } from '../../lib/Tldraw'
+import { TLUiOverrides } from '../../lib/ui/overrides'
+import {
+	renderTldrawComponent,
+	renderTldrawComponentWithEditor,
+} from '../testutils/renderTldrawComponent'
+
+it('opens on right-click', async () => {
+	await renderTldrawComponent(
+		<Tldraw
+			onMount={(editor) => {
+				editor.createShape({ id: createShapeId(), type: 'geo' })
+			}}
+		/>
+	)
+	const canvas = await screen.findByTestId('canvas')
+
+	fireEvent.contextMenu(canvas)
+	await screen.findByTestId('context-menu')
+	await screen.findByTestId('menu-item.select-all')
+
+	fireEvent.keyDown(canvas, { key: 'Escape' })
+	expect(screen.queryByTestId('context-menu')).toBeNull()
+})
+
+it('updates overrides reactively', async () => {
+	const overrides: TLUiOverrides = {
+		contextMenu: (editor, schema) => {
+			const items = editor.getSelectedShapeIds().length
+			if (items === 0) return schema
+			return [
+				...schema,
+				{
+					type: 'item',
+					id: 'tester',
+					disabled: false,
+					readonlyOk: true,
+					checked: false,
+					actionItem: {
+						id: 'tester',
+						readonlyOk: true,
+						onSelect: noop,
+						label: `Selected: ${items}`,
+					},
+				},
+			]
+		},
+	}
+	const { editor } = await renderTldrawComponentWithEditor((onMount) => (
+		<Tldraw onMount={onMount} overrides={overrides} />
+	))
+
+	await act(() => editor.createShape({ id: createShapeId(), type: 'geo' }).selectAll())
+
+	// open the context menu:
+	fireEvent.contextMenu(await screen.findByTestId('canvas'))
+
+	// check that the context menu item was added:
+	await screen.findByTestId('menu-item.tester')
+
+	// It should disappear when we deselect all shapes:
+	await act(() => editor.setSelectedShapes([]))
+	expect(screen.queryByTestId('menu-item.tester')).toBeNull()
+
+	// It should update its label when it changes:
+	await act(() => editor.selectAll())
+	const item = await screen.findByTestId('menu-item.tester')
+	expect(item.textContent).toBe('Selected: 1')
+	await act(() => editor.createShape({ id: createShapeId(), type: 'geo' }).selectAll())
+	expect(item.textContent).toBe('Selected: 2')
+})