UnitControl: Fix internal unit state/parsing to handle incoming unit …

…prop (#23521)

This update fixes the `<UnitControl />` component to handle (new) incoming `unit` prop changes. This is important for handling interactions such as unit swapping - an interaction we have for the Cover block.

packages/components/src/unit-control/index.js

@@ -7,7 +7,7 @@ import classnames from 'classnames';
 /**
  * WordPress dependencies
  */
-import { useState, forwardRef } from '@wordpress/element';
+import { forwardRef } from '@wordpress/element';
 
 /**
  * Internal dependencies
@@ -19,6 +19,7 @@ import {
 import { Root, ValueInput } from './styles/unit-control-styles';
 import UnitSelectControl from './unit-select-control';
 import { CSS_UNITS, getParsedValue, getValidParsedUnit } from './utils';
+import { useControlledState } from '../utils/hooks';
 
 function UnitControl(
 	{
@@ -43,7 +44,9 @@ function UnitControl(
 	ref
 ) {
 	const [ value, initialUnit ] = getParsedValue( valueProp, unitProp, units );
-	const [ unit, setUnit ] = useState( initialUnit );
+	const [ unit, setUnit ] = useControlledState( unitProp, {
+		initial: initialUnit,
+	} );
 
 	const classes = classnames( 'components-unit-control', className );
 

packages/components/src/unit-control/test/index.js

@@ -3,6 +3,7 @@
  */
 import { render, unmountComponentAtNode } from 'react-dom';
 import { act, Simulate } from 'react-dom/test-utils';
+import { render as testRender } from '@testing-library/react';
 
 /**
  * WordPress dependencies
@@ -372,5 +373,19 @@ describe( 'UnitControl', () => {
 
 			expect( state ).toBe( '123rem' );
 		} );
+
+		it( 'should update unit after initial render and with new unit prop', () => {
+			const { container: testContainer, rerender } = testRender(
+				<UnitControl value={ '10%' } />
+			);
+
+			const select = testContainer.querySelector( 'select' );
+
+			expect( select.value ).toBe( '%' );
+
+			rerender( <UnitControl value={ state } unit="em" /> );
+
+			expect( select.value ).toBe( 'em' );
+		} );
 	} );
 } );