Fix DisableNonPageContentBlocks

- Updates DisableNonPageContentBlocks to only affect top level content.
  For example, a Featured Image within a Content block should remain as
  is.

- Refactors DisableNonPageContentBlocks to use a selector so that we can
  write unit tests for the logic.

packages/editor/src/components/provider/disable-non-page-content-blocks.js

@@ -2,54 +2,40 @@
  * WordPress dependencies
  */
 import { useSelect, useDispatch } from '@wordpress/data';
-import {
-	useBlockEditingMode,
-	store as blockEditorStore,
-} from '@wordpress/block-editor';
+import { store as blockEditorStore } from '@wordpress/block-editor';
 import { useEffect } from '@wordpress/element';
 
 /**
  * Internal dependencies
  */
-import { PAGE_CONTENT_BLOCK_TYPES } from './constants';
+import { unlock } from '../../lock-unlock';
+import { store as editorStore } from '../../store';
 
-function DisableBlock( { clientId } ) {
-	const isDescendentOfQueryLoop = useSelect(
-		( select ) => {
-			const { getBlockParentsByBlockName } = select( blockEditorStore );
-			return (
-				getBlockParentsByBlockName( clientId, 'core/query' ).length !==
-				0
-			);
-		},
-		[ clientId ]
+function useDisableNonPageContentBlocks() {
+	const contentIds = useSelect(
+		( select ) => unlock( select( editorStore ) ).getPageContentBlocks(),
+		[]
 	);
-	const mode = isDescendentOfQueryLoop ? undefined : 'contentOnly';
 	const { setBlockEditingMode, unsetBlockEditingMode } =
 		useDispatch( blockEditorStore );
 	useEffect( () => {
-		if ( mode ) {
-			setBlockEditingMode( clientId, mode );
-			return () => {
-				unsetBlockEditingMode( clientId );
-			};
+		setBlockEditingMode( '', 'disabled' ); // Disable editing at the root level.
+		for ( const contentId of contentIds ) {
+			setBlockEditingMode( contentId, 'contentOnly' ); // Re-enable each content block.
 		}
-	}, [ clientId, mode, setBlockEditingMode, unsetBlockEditingMode ] );
+		return () => {
+			unsetBlockEditingMode( '' );
+			for ( const contentId of contentIds ) {
+				unsetBlockEditingMode( contentId );
+			}
+		};
+	}, [ contentIds, setBlockEditingMode, unsetBlockEditingMode ] );
 }
 
 /**
  * Component that when rendered, makes it so that the site editor allows only
  * page content to be edited.
  */
 export default function DisableNonPageContentBlocks() {
-	useBlockEditingMode( 'disabled' );
-	const clientIds = useSelect( ( select ) => {
-		return select( blockEditorStore ).getBlocksByName(
-			PAGE_CONTENT_BLOCK_TYPES
-		);
-	}, [] );
-
-	return clientIds.map( ( clientId ) => {
-		return <DisableBlock key={ clientId } clientId={ clientId } />;
-	} );
+	useDisableNonPageContentBlocks();
 }

packages/editor/src/store/private-selectors.js

@@ -49,3 +49,26 @@ export const getInsertionPoint = createRegistrySelector(
 export function getListViewToggleRef( state ) {
 	return state.listViewToggleRef;
 }
+
+export const getPageContentBlocks = createRegistrySelector(
+	( select ) =>
+		( state, rootClientId = '' ) => {
+			const { getBlockOrder, getBlockName } = select( blockEditorStore );
+			const contentIds = [];
+			for ( const clientId of getBlockOrder( rootClientId ) ) {
+				const blockName = getBlockName( clientId );
+				if (
+					blockName === 'core/post-title' ||
+					blockName === 'core/post-featured-image' ||
+					blockName === 'core/post-content'
+				) {
+					contentIds.push( clientId );
+				} else if ( blockName !== 'core/query' ) {
+					contentIds.push(
+						...getPageContentBlocks( state, clientId )
+					);
+				}
+			}
+			return contentIds;
+		}
+);

packages/editor/src/store/test/private-selectors.js

@@ -0,0 +1,88 @@
+/**
+ * Internal dependencies
+ */
+import { getPageContentBlocks } from '../private-selectors';
+
+describe( 'private selectors', () => {
+	describe( 'getPageContentBlocks', () => {
+		// |-0 template-part
+		// | |- 00 site-title
+		// | |- 01 navigation
+		// |-1 group
+		// | |-10 post-title
+		// | |-11 post-featured-image
+		// | |-12 post-content
+		// | | |-120 paragraph
+		// | | |-121 post-featured-image
+		// |-2 query
+		// | |-20 post-title
+		// | |-21 post-featured-image
+		// | |-22 post-content
+		// |-3 template-part
+		// | |-30 paragraph
+
+		const testBlocks = {
+			'': { order: [ '0', '1', '2' ] },
+			0: {
+				name: 'core/template-part',
+				order: [ '00', '01' ],
+			},
+			'00': { name: 'core/site-title' },
+			'01': { name: 'core/navigation' },
+			1: {
+				name: 'core/group',
+				order: [ '10', '11', '12' ],
+			},
+			10: { name: 'core/post-title' },
+			11: { name: 'core/post-featured-image' },
+			12: {
+				name: 'core/post-content',
+				order: [ '120', '121' ],
+			},
+			120: { name: 'core/paragraph' },
+			121: { name: 'core/post-featured-image' },
+			2: {
+				name: 'core/query',
+				order: [ '20', '21', '22' ],
+			},
+			20: { name: 'core/post-title' },
+			21: { name: 'core/post-featured-image' },
+			22: { name: 'core/post-content' },
+			3: {
+				name: 'core/template-part',
+				order: [ '30' ],
+			},
+			30: { name: 'core/paragraph' },
+		};
+
+		const getBlockOrder = jest.fn();
+		const getBlockName = jest.fn();
+		getPageContentBlocks.registry = {
+			select: jest.fn( () => ( {
+				getBlockOrder,
+				getBlockName,
+			} ) ),
+		};
+
+		it( 'returns an empty array if there are no blocks', () => {
+			getBlockOrder.mockReturnValueOnce( [] );
+			expect( getPageContentBlocks( {} ) ).toEqual( [] );
+		} );
+
+		it( 'returns page content blocks', () => {
+			getBlockOrder.mockImplementation(
+				( rootClientId ) => testBlocks[ rootClientId ]?.order ?? []
+			);
+			getBlockName.mockImplementation(
+				( clientId ) => testBlocks[ clientId ]?.name
+			);
+			expect( getPageContentBlocks( {} ) ).toEqual( [
+				'10', // post-title
+				'11', // post-featured-image
+				'12', // post-content
+				// NOT the post-featured-image nested within post-content
+				// NOT any of the content blocks within query
+			] );
+		} );
+	} );
+} );