diff --git a/packages/components/src/unit-control/index.js b/packages/components/src/unit-control/index.js
index 39137c6c132a7..06f3f9a84d5c6 100644
--- a/packages/components/src/unit-control/index.js
+++ b/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 );
diff --git a/packages/components/src/unit-control/test/index.js b/packages/components/src/unit-control/test/index.js
index c6b0431c96847..f78abf47781c6 100644
--- a/packages/components/src/unit-control/test/index.js
+++ b/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(
+
+ );
+
+ const select = testContainer.querySelector( 'select' );
+
+ expect( select.value ).toBe( '%' );
+
+ rerender( );
+
+ expect( select.value ).toBe( 'em' );
+ } );
} );
} );