Use useSelect() approach instead of private selector

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

@@ -5,21 +5,33 @@ import { useSelect, useDispatch } from '@wordpress/data';
 import { store as blockEditorStore } from '@wordpress/block-editor';
 import { useEffect } from '@wordpress/element';
 
-/**
- * Internal dependencies
- */
-import { unlock } from '../../lock-unlock';
-import { store as editorStore } from '../../store';
+const PAGE_CONTENT_BLOCKS = [
+	'core/post-title',
+	'core/post-featured-image',
+	'core/post-content',
+];
 
 function useDisableNonPageContentBlocks() {
-	const contentIds = useSelect(
-		( select ) => unlock( select( editorStore ) ).getPageContentBlocks(),
-		[]
-	);
+	const contentIds = useSelect( ( select ) => {
+		const { getBlocksByName, getBlockParents, getBlockName } =
+			select( blockEditorStore );
+		return getBlocksByName( PAGE_CONTENT_BLOCKS ).filter( ( clientId ) =>
+			getBlockParents( clientId ).every( ( parentClientId ) => {
+				const parentBlockName = getBlockName( parentClientId );
+				return (
+					parentBlockName !== 'core/query' &&
+					! PAGE_CONTENT_BLOCKS.includes( parentBlockName )
+				);
+			} )
+		);
+	}, [] );
+
 	const { setBlockEditingMode, unsetBlockEditingMode } =
 		useDispatch( blockEditorStore );
+
 	useEffect( () => {
 		setBlockEditingMode( '', 'disabled' ); // Disable editing at the root level.
+
 		for ( const contentId of contentIds ) {
 			setBlockEditingMode( contentId, 'contentOnly' ); // Re-enable each content block.
 		}

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

@@ -0,0 +1,90 @@
+/**
+ * External dependencies
+ */
+import { render } from '@testing-library/react';
+
+/**
+ * WordPress dependencies
+ */
+import { createRegistry, RegistryProvider } from '@wordpress/data';
+
+/**
+ * Internal dependencies
+ */
+import DisableNonPageContentBlocks from '../disable-non-page-content-blocks';
+
+describe( 'DisableNonPageContentBlocks', () => {
+	it( 'disables page content blocks', () => {
+		const testBlocks = {
+			0: 'core/template-part',
+			/**/ '00': 'core/site-title',
+			/**/ '01': 'core/navigation',
+			1: 'core/group',
+			/**/ 10: 'core/post-title',
+			/**/ 11: 'core/post-featured-image',
+			/**/ 12: 'core/post-content',
+			/**/ /**/ 120: 'core/paragraph',
+			/**/ /**/ 121: 'core/post-featured-image',
+			2: 'core/query',
+			/**/ 20: 'core/post-title',
+			/**/ 21: 'core/post-featured-image',
+			/**/ 22: 'core/post-content',
+			3: 'core/template-part',
+			/**/ 30: 'core/paragraph',
+		};
+
+		const setBlockEditingMode = jest.fn( () => ( {
+			type: 'SET_BLOCK_EDITING_MODE',
+		} ) );
+		const unsetBlockEditingMode = jest.fn( () => ( {
+			type: 'UNSET_BLOCK_EDITING_MODE',
+		} ) );
+
+		const registry = createRegistry( {
+			'core/block-editor': {
+				reducer: () => {},
+				selectors: {
+					getBlocksByName( state, blockNames ) {
+						return Object.keys( testBlocks ).filter( ( clientId ) =>
+							blockNames.includes( testBlocks[ clientId ] )
+						);
+					},
+					getBlockParents( state, clientId ) {
+						return clientId.slice( 0, -1 ).split( '' );
+					},
+					getBlockName( state, clientId ) {
+						return testBlocks[ clientId ];
+					},
+				},
+				actions: {
+					setBlockEditingMode,
+					unsetBlockEditingMode,
+				},
+			},
+		} );
+
+		const { unmount } = render(
+			<RegistryProvider value={ registry }>
+				<DisableNonPageContentBlocks />
+			</RegistryProvider>
+		);
+
+		expect( setBlockEditingMode.mock.calls ).toEqual( [
+			[ '', 'disabled' ], // root
+			[ '10', 'contentOnly' ], // post-title
+			[ '11', 'contentOnly' ], // post-featured-image
+			[ '12', 'contentOnly' ], // post-content
+			// NOT the post-featured-image nested within post-content
+			// NOT any of the content blocks within query
+		] );
+
+		unmount();
+
+		expect( unsetBlockEditingMode.mock.calls ).toEqual( [
+			[ '' ], // root
+			[ '10' ], // post-title
+			[ '11' ], // post-featured-image
+			[ '12' ], // post-content
+		] );
+	} );
+} );

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

@@ -49,26 +49,3 @@ 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;
-		}
-);