From a468cf59180a694a5da104df3fa4aca23f0f6eac Mon Sep 17 00:00:00 2001
From: Andrew Duthie <andrew@andrewduthie.com>
Date: Wed, 7 Mar 2018 11:02:23 -0500
Subject: [PATCH] Data: Allow undefined return from withSelect mapSelectToProps
 (#5421)

* Data: Avoid lingering props on subsequent mapping changes

* Data: Allow undefined return from mapSelectToProps
---
 data/index.js      | 12 ++++++++----
 data/test/index.js | 48 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/data/index.js b/data/index.js
index 765a3d52cd3d2d..bf3558b0da5347 100644
--- a/data/index.js
+++ b/data/index.js
@@ -218,14 +218,18 @@ export const withSelect = ( mapStateToProps ) => ( WrappedComponent ) => {
 				return;
 			}
 
-			const newState = mapStateToProps( select, props );
-			if ( ! isEqualShallow( newState, this.state ) ) {
-				this.setState( newState );
+			const { mergeProps } = this.state;
+			const nextMergeProps = mapStateToProps( select, props ) || {};
+
+			if ( ! isEqualShallow( nextMergeProps, mergeProps ) ) {
+				this.setState( {
+					mergeProps: nextMergeProps,
+				} );
 			}
 		}
 
 		render() {
-			return <WrappedComponent { ...this.props } { ...this.state } />;
+			return <WrappedComponent { ...this.props } { ...this.state.mergeProps } />;
 		}
 	}
 
diff --git a/data/test/index.js b/data/test/index.js
index e638e2ba1ccf74..1e03d0f546950f 100644
--- a/data/test/index.js
+++ b/data/test/index.js
@@ -273,6 +273,54 @@ describe( 'withSelect', () => {
 
 			expect( mapSelectToProps ).toHaveBeenCalledTimes( 1 );
 		} );
+
+		itWithExtraAssertions( 'omits props which are not returned on subsequent mappings', () => {
+			registerReducer( 'demo', ( state = 'OK' ) => state );
+			registerSelectors( 'demo', {
+				getValue: ( state ) => state,
+			} );
+
+			const Component = withSelectImplementation( ( _select, ownProps ) => {
+				return {
+					[ ownProps.propName ]: _select( 'demo' ).getValue(),
+				};
+			} )( () => <div /> );
+
+			wrapper = mount( <Component propName="foo" /> );
+
+			expect( wrapper.childAt( 0 ).props() ).toEqual( { foo: 'OK', propName: 'foo' } );
+
+			wrapper.setProps( { propName: 'bar' } );
+
+			expect( wrapper.childAt( 0 ).props() ).toEqual( { bar: 'OK', propName: 'bar' } );
+		} );
+
+		itWithExtraAssertions( 'allows undefined return from mapSelectToProps', () => {
+			registerReducer( 'demo', ( state = 'OK' ) => state );
+			registerSelectors( 'demo', {
+				getValue: ( state ) => state,
+			} );
+
+			const Component = withSelectImplementation( ( _select, ownProps ) => {
+				if ( ownProps.pass ) {
+					return {
+						count: _select( 'demo' ).getValue(),
+					};
+				}
+			} )( ( props ) => <div>{ props.count || 'Unknown' }</div> );
+
+			wrapper = mount( <Component pass={ false } /> );
+
+			expect( wrapper.childAt( 0 ).text() ).toBe( 'Unknown' );
+
+			wrapper.setProps( { pass: true } );
+
+			expect( wrapper.childAt( 0 ).text() ).toBe( 'OK' );
+
+			wrapper.setProps( { pass: false } );
+
+			expect( wrapper.childAt( 0 ).text() ).toBe( 'Unknown' );
+		} );
 	}
 
 	cases( withSelect );