From 55c12fe6c52ae003aac37a835de1934b989f203f Mon Sep 17 00:00:00 2001
From: Jefferson Rabb <jeff@atavist.net>
Date: Wed, 22 Apr 2020 23:54:43 -0400
Subject: [PATCH] fix: regression of test email error/success notifications

---
 src/components/with-api-handler/index.js |  4 ++++
 src/editor/testing/index.js              | 14 +++++++++-----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/components/with-api-handler/index.js b/src/components/with-api-handler/index.js
index f429cf50c..f5d80b0f8 100644
--- a/src/components/with-api-handler/index.js
+++ b/src/components/with-api-handler/index.js
@@ -12,6 +12,9 @@ export default () =>
 			const [ inFlight, setInFlight ] = useState( false );
 			const { createSuccessNotice, createErrorNotice, removeNotice } = dispatch( 'core/notices' );
 			const { getNotices } = select( 'core/notices' );
+			const setInFlightForAsync = () => {
+				setInFlight( true );
+			};
 			const apiFetchWithErrorHandling = apiRequest => {
 				setInFlight( true );
 				return new Promise( ( resolve, reject ) => {
@@ -38,6 +41,7 @@ export default () =>
 				<OriginalComponent
 					{ ...props }
 					apiFetchWithErrorHandling={ apiFetchWithErrorHandling }
+					setInFlightForAsync={ setInFlightForAsync }
 					inFlight={ inFlight }
 				/>
 			);
diff --git a/src/editor/testing/index.js b/src/editor/testing/index.js
index 92e7f7dc5..6e807656a 100644
--- a/src/editor/testing/index.js
+++ b/src/editor/testing/index.js
@@ -2,14 +2,19 @@
  * WordPress dependencies
  */
 import { __ } from '@wordpress/i18n';
-import apiFetch from '@wordpress/api-fetch';
 import { withSelect, withDispatch } from '@wordpress/data';
 import { compose } from '@wordpress/compose';
 import { Fragment, useState } from '@wordpress/element';
 import { Button, TextControl } from '@wordpress/components';
 import { hasValidEmail } from '../utils';
 
+/**
+ * Internal dependencies
+ */
+import withApiHandler from '../../components/with-api-handler';
+
 export default compose( [
+	withApiHandler(),
 	withSelect( select => {
 		const { getCurrentPostId } = select( 'core/editor' );
 		return { postId: getCurrentPostId() };
@@ -20,11 +25,10 @@ export default compose( [
 			savePost,
 		};
 	} ),
-] )( ( { postId, savePost } ) => {
-	const [ inFlight, setInFlight ] = useState( false );
+] )( ( { apiFetchWithErrorHandling, inFlight, postId, savePost, setInFlightForAsync } ) => {
 	const [ testEmail, setTestEmail ] = useState( '' );
 	const sendMailchimpTest = async () => {
-		setInFlight( true );
+		setInFlightForAsync();
 		await savePost();
 		const params = {
 			path: `/newspack-newsletters/v1/mailchimp/${ postId }/test`,
@@ -33,7 +37,7 @@ export default compose( [
 			},
 			method: 'POST',
 		};
-		apiFetch( params ).then( () => setInFlight( false ) );
+		apiFetchWithErrorHandling( params );
 	};
 	return (
 		<Fragment>