Fix 'Country is required' error on the Cart block when updating shipping address

assets/js/base/context/hooks/cart/use-store-cart.ts

@@ -184,6 +184,7 @@ export const useStoreCart = (
 			const cartIsLoading = ! store.hasFinishedResolution(
 				'getCartData'
 			);
+
 			const shippingRatesLoading = store.isCustomerDataUpdating();
 			const { receiveCart } = dispatch( storeKey );
 			const billingAddress = decodeValues( cartData.billingAddress );

assets/js/base/context/hooks/use-customer-data.ts

@@ -11,9 +11,13 @@ import {
 	pluckAddress,
 	pluckEmail,
 } from '@woocommerce/base-utils';
-import type {
+import {
 	CartResponseBillingAddress,
 	CartResponseShippingAddress,
+	objectHasProp,
+	BillingAddressShippingAddress,
+	CartBillingAddress,
+	CartShippingAddress,
 } from '@woocommerce/types';
 
 declare type CustomerData = {
@@ -86,12 +90,40 @@ export const useCustomerData = (): {
 		billingAddress: CartResponseBillingAddress;
 	} = useStoreCart();
 
+	// We need to store a ref to the shipping and billing addresses in order to re-populate the local state
+	// here with the hydrated values once they change. The refs keep track of the changing values of the addresses
+	// as they are hydrated.
+	const billingDataRef = useRef( initialBillingAddress );
+	const shippingAddressRef = useRef( initialShippingAddress );
+	billingDataRef.current = initialBillingAddress;
+	shippingAddressRef.current = initialShippingAddress;
+
 	// State of customer data is tracked here from this point, using the initial values from the useStoreCart hook.
 	const [ customerData, setCustomerData ] = useState< CustomerData >( {
 		billingData: initialBillingAddress,
 		shippingAddress: initialShippingAddress,
 	} );
 
+	// We only want to update the local state once, otherwise the data on the checkout page gets overwritten
+	// with the initial state of the addresses here
+	const [ hasCustomerDataSynced, setHasCustomerDataSynced ] = useState<
+		boolean
+	>( false );
+
+	if (
+		! hasCustomerDataSynced &&
+		shouldUpdateAddressStore(
+			customerData.shippingAddress,
+			shippingAddressRef.current
+		)
+	) {
+		setCustomerData( {
+			billingData: billingDataRef.current,
+			shippingAddress: shippingAddressRef.current,
+		} );
+		setHasCustomerDataSynced( true );
+	}
+
 	// Store values last sent to the server in a ref to avoid requests unless important fields are changed.
 	const previousCustomerData = useRef< CustomerData >( customerData );
 
@@ -146,23 +178,46 @@ export const useCustomerData = (): {
 	 */
 	useEffect( () => {
 		// Only push updates when enough fields are populated.
+		const shouldUpdateBillingAddress = shouldUpdateAddressStore(
+			previousCustomerData.current.billingData,
+			debouncedCustomerData.billingData
+		);
+
+		const shouldUpdateShippingAddress = shouldUpdateAddressStore(
+			previousCustomerData.current.shippingAddress,
+			debouncedCustomerData.shippingAddress
+		);
+
+		if ( ! shouldUpdateBillingAddress && ! shouldUpdateShippingAddress ) {
+			return;
+		}
+
+		const customerDataToUpdate:
+			| Partial< BillingAddressShippingAddress >
+			| Record<
+					keyof BillingAddressShippingAddress,
+					CartBillingAddress | CartShippingAddress
+			  > = {};
+
+		if ( shouldUpdateBillingAddress ) {
+			customerDataToUpdate.billing_address =
+				debouncedCustomerData.billingData;
+		}
+		if ( shouldUpdateShippingAddress ) {
+			customerDataToUpdate.shipping_address =
+				debouncedCustomerData.shippingAddress;
+		}
+
 		if (
-			! shouldUpdateAddressStore(
-				previousCustomerData.current.billingData,
-				debouncedCustomerData.billingData
-			) &&
-			! shouldUpdateAddressStore(
-				previousCustomerData.current.shippingAddress,
-				debouncedCustomerData.shippingAddress
-			)
+			! objectHasProp( customerDataToUpdate, 'billing_address' ) &&
+			! objectHasProp( customerDataToUpdate, 'shipping_address' )
 		) {
 			return;
 		}
 		previousCustomerData.current = debouncedCustomerData;
-		updateCustomerData( {
-			billing_address: debouncedCustomerData.billingData,
-			shipping_address: debouncedCustomerData.shippingAddress,
-		} )
+		updateCustomerData(
+			customerDataToUpdate as Partial< BillingAddressShippingAddress >
+		)
 			.then( () => {
 				removeNotice( 'checkout' );
 			} )

assets/js/data/cart/actions.ts

@@ -6,9 +6,8 @@ import type {
 	Cart,
 	CartResponse,
 	CartResponseItem,
-	CartBillingAddress,
-	CartShippingAddress,
 	ExtensionCartUpdateArgs,
+	BillingAddressShippingAddress,
 } from '@woocommerce/types';
 import { ReturnOrGeneratorYieldUnion } from '@automattic/data-stores';
 import { camelCase, mapKeys } from 'lodash';
@@ -463,11 +462,6 @@ export function* selectShippingRate(
 	return true;
 }
 
-type BillingAddressShippingAddress = {
-	billing_address: CartBillingAddress;
-	shipping_address: CartShippingAddress;
-};
-
 /**
  * Updates the shipping and/or billing address for the customer and returns an
  * updated cart.
@@ -476,7 +470,7 @@ type BillingAddressShippingAddress = {
  *   billing_address and shipping_address.
  */
 export function* updateCustomerData(
-	customerData: BillingAddressShippingAddress
+	customerData: Partial< BillingAddressShippingAddress >
 ): Generator< unknown, boolean, { response: CartResponse } > {
 	yield updatingCustomerData( true );
 

assets/js/types/type-defs/cart.ts

@@ -198,3 +198,8 @@ export interface ExtensionCartUpdateArgs {
 	data: Record< string, unknown >;
 	namespace: string;
 }
+
+export interface BillingAddressShippingAddress {
+	billing_address: CartBillingAddress;
+	shipping_address: CartShippingAddress;
+}