Fix/extension data conflicts

assets/js/base/context/hooks/use-checkout-extension-data.ts

@@ -2,8 +2,7 @@
  * External dependencies
  */
 import { useDispatch, useSelect } from '@wordpress/data';
-import { useCallback, useEffect, useRef } from '@wordpress/element';
-import isShallowEqual from '@wordpress/is-shallow-equal';
+import { useCallback, useRef } from '@wordpress/element';
 import { CHECKOUT_STORE_KEY } from '@woocommerce/block-data';
 
 /**
@@ -28,28 +27,17 @@ export const useCheckoutExtensionData = (): {
 	);
 	const extensionDataRef = useRef( extensionData );
 
-	useEffect( () => {
-		if ( ! isShallowEqual( extensionData, extensionDataRef.current ) ) {
-			extensionDataRef.current = extensionData;
-		}
-	}, [ extensionData ] );
-
-	const setExtensionDataWithNamespace = useCallback(
+	const setExtensionData = useCallback(
 		( namespace, key, value ) => {
-			const currentData = extensionDataRef.current[ namespace ] || {};
-			__internalSetExtensionData( {
-				...extensionDataRef.current,
-				[ namespace ]: {
-					...currentData,
-					[ key ]: value,
-				},
+			__internalSetExtensionData( namespace, {
+				[ key ]: value,
 			} );
 		},
 		[ __internalSetExtensionData ]
 	);
 
 	return {
 		extensionData: extensionDataRef.current,
-		setExtensionData: setExtensionDataWithNamespace,
+		setExtensionData,
 	};
 };

assets/js/data/checkout/actions.ts

@@ -131,15 +131,20 @@ export const setPrefersCollection = ( prefersCollection: boolean ) => ( {
 } );
 
 /**
- * Register some extra data for an extension. This works with the
- *
- * @param  extensionData An object containing the data to register for an extension
+ * Registers additional data under an extension namespace.
  */
 export const __internalSetExtensionData = (
-	extensionData: Record< string, Record< string, unknown > >
+	// The namespace for the extension. Defaults to 'default'. Must be unique to prevent conflicts.
+	namespace: string,
+	// Data to register under the namespace.
+	extensionData: Record< string, unknown >,
+	// If true, all data under the current extension namespace is replaced. If false, data is appended.
+	replace = false
 ) => ( {
 	type: types.SET_EXTENSION_DATA,
 	extensionData,
+	namespace,
+	replace,
 } );
 
 export type CheckoutAction =

assets/js/data/checkout/reducers.ts

@@ -149,11 +149,19 @@ const reducer = ( state = defaultState, action: CheckoutAction ) => {
 		case types.SET_EXTENSION_DATA:
 			if (
 				action.extensionData !== undefined &&
-				state.extensionData !== action.extensionData
+				action.namespace !== undefined
 			) {
 				newState = {
 					...state,
-					extensionData: action.extensionData,
+					extensionData: {
+						...state.extensionData,
+						[ action.namespace ]: action.replace
+							? action.extensionData
+							: {
+									...state.extensionData[ action.namespace ],
+									...action.extensionData,
+							  },
+					},
 				};
 			}
 			break;

assets/js/data/checkout/test/reducer.ts

@@ -215,21 +215,81 @@ describe.only( 'Checkout Store Reducer', () => {
 		).toEqual( expectedState );
 	} );
 
-	it( 'should handle SET_EXTENSION_DATA', () => {
-		const mockExtensionData = {
-			testExtension: { key: 'test', value: 'test2' },
-		};
-		const expectedState = {
-			...defaultState,
-			status: STATUS.IDLE,
-			extensionData: mockExtensionData,
-		};
-
-		expect(
-			reducer(
+	describe( 'should handle SET_EXTENSION_DATA', () => {
+		it( 'should set data under a namespace', () => {
+			const mockExtensionData = {
+				extensionNamespace: {
+					testKey: 'test-value',
+					testKey2: 'test-value-2',
+				},
+			};
+			const expectedState = {
+				...defaultState,
+				status: STATUS.IDLE,
+				extensionData: mockExtensionData,
+			};
+			expect(
+				reducer(
+					defaultState,
+					actions.__internalSetExtensionData(
+						'extensionNamespace',
+						mockExtensionData.extensionNamespace
+					)
+				)
+			).toEqual( expectedState );
+		} );
+		it( 'should append data under a namespace', () => {
+			const mockExtensionData = {
+				extensionNamespace: {
+					testKey: 'test-value',
+					testKey2: 'test-value-2',
+				},
+			};
+			const expectedState = {
+				...defaultState,
+				status: STATUS.IDLE,
+				extensionData: mockExtensionData,
+			};
+			const firstState = reducer(
 				defaultState,
-				actions.__internalSetExtensionData( mockExtensionData )
-			)
-		).toEqual( expectedState );
+				actions.__internalSetExtensionData( 'extensionNamespace', {
+					testKey: 'test-value',
+				} )
+			);
+			const secondState = reducer(
+				firstState,
+				actions.__internalSetExtensionData( 'extensionNamespace', {
+					testKey2: 'test-value-2',
+				} )
+			);
+			expect( secondState ).toEqual( expectedState );
+		} );
+		it( 'support replacing data under a namespace', () => {
+			const mockExtensionData = {
+				extensionNamespace: {
+					testKey: 'test-value',
+				},
+			};
+			const expectedState = {
+				...defaultState,
+				status: STATUS.IDLE,
+				extensionData: mockExtensionData,
+			};
+			const firstState = reducer(
+				defaultState,
+				actions.__internalSetExtensionData( 'extensionNamespace', {
+					testKeyOld: 'test-value',
+				} )
+			);
+			const secondState = reducer(
+				firstState,
+				actions.__internalSetExtensionData(
+					'extensionNamespace',
+					{ testKey: 'test-value' },
+					true
+				)
+			);
+			expect( secondState ).toEqual( expectedState );
+		} );
 	} );
 } );

docs/internal-developers/block-client-apis/checkout/checkout-api.md

@@ -71,7 +71,7 @@ The following actions can be dispatched from the Checkout data store:
 -   `__internalSetUseShippingAsBilling( useShippingAsBilling: boolean )`: Set `state.useShippingAsBilling` to `useShippingAsBilling`
 -   `__internalSetShouldCreateAccount( shouldCreateAccount: boolean )`: Set `state.shouldCreateAccount` to `shouldCreateAccount`
 -   `__internalSetOrderNotes( orderNotes: string )`: Set `state.orderNotes` to `orderNotes`
--   `__internalSetExtensionData( extensionData: Record< string, Record< string, unknown > > )`: Set `state.extensionData` to `extensionData`
+-   `__internalSetExtensionData( namespace: string, extensionData: Record< string, unknown > )`: Set `state.extensionData` to `extensionData`
 
 ### Contexts