From 306d3422edf41b3fd960c950946c38ba9aa01063 Mon Sep 17 00:00:00 2001
From: Jon Q <hello@jonquach.com>
Date: Fri, 26 Jun 2020 13:19:56 -0400
Subject: [PATCH] UnitControl: Fix internal unit state/parsing to handle
 incoming unit prop

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 +++++--
 .../components/src/unit-control/test/index.js     | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

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(
+				<UnitControl value={ '10%' } />
+			);
+
+			const select = testContainer.querySelector( 'select' );
+
+			expect( select.value ).toBe( '%' );
+
+			rerender( <UnitControl value={ state } unit="em" /> );
+
+			expect( select.value ).toBe( 'em' );
+		} );
 	} );
 } );