From f50d8fcff2a45594553e14303c4872778e4adc25 Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Tue, 29 Jan 2019 17:39:24 -0800
Subject: [PATCH] Improve Reducer Hook's lazy init API (#14723)

* Improve Reducer Hook's lazy init API

* Use generic type for initilizer input

Still requires an `any` cast in the case where `init` function is
not provided.
---
 src/ReactShallowRenderer.js                   | 19 ++++++++------
 .../ReactShallowRendererHooks-test.js         | 26 ++++++-------------
 2 files changed, 19 insertions(+), 26 deletions(-)

diff --git a/src/ReactShallowRenderer.js b/src/ReactShallowRenderer.js
index ee8b68b..0701854 100644
--- a/src/ReactShallowRenderer.js
+++ b/src/ReactShallowRenderer.js
@@ -223,10 +223,10 @@ class ReactShallowRenderer {
   }
 
   _createDispatcher(): DispatcherType {
-    const useReducer = <S, A>(
+    const useReducer = <S, I, A>(
       reducer: (S, A) => S,
-      initialState: S,
-      initialAction: A | void | null,
+      initialArg: I,
+      init?: I => S,
     ): [S, Dispatch<A>] => {
       this._validateCurrentlyRenderingComponent();
       this._createWorkInProgressHook();
@@ -259,13 +259,16 @@ class ReactShallowRenderer {
         }
         return [workInProgressHook.memoizedState, dispatch];
       } else {
+        let initialState;
         if (reducer === basicStateReducer) {
           // Special case for `useState`.
-          if (typeof initialState === 'function') {
-            initialState = initialState();
-          }
-        } else if (initialAction !== undefined && initialAction !== null) {
-          initialState = reducer(initialState, initialAction);
+          initialState =
+            typeof initialArg === 'function'
+              ? ((initialArg: any): () => S)()
+              : ((initialArg: any): S);
+        } else {
+          initialState =
+            init !== undefined ? init(initialArg) : ((initialArg: any): S);
         }
         workInProgressHook.memoizedState = initialState;
         const queue: UpdateQueue<A> = (workInProgressHook.queue = {
diff --git a/src/__tests__/ReactShallowRendererHooks-test.js b/src/__tests__/ReactShallowRendererHooks-test.js
index 60c79a1..da352a9 100644
--- a/src/__tests__/ReactShallowRendererHooks-test.js
+++ b/src/__tests__/ReactShallowRendererHooks-test.js
@@ -91,23 +91,19 @@ describe('ReactShallowRenderer with hooks', () => {
   });
 
   it('should work with useReducer', () => {
-    const initialState = {count: 0};
-
     function reducer(state, action) {
       switch (action.type) {
-        case 'reset':
-          return initialState;
         case 'increment':
           return {count: state.count + 1};
         case 'decrement':
           return {count: state.count - 1};
-        default:
-          return state;
       }
     }
 
-    function SomeComponent({initialCount}) {
-      const [state] = React.useReducer(reducer, {count: initialCount});
+    function SomeComponent(props) {
+      const [state] = React.useReducer(reducer, props, p => ({
+        count: p.initialCount,
+      }));
 
       return (
         <div>
@@ -141,25 +137,19 @@ describe('ReactShallowRenderer with hooks', () => {
   });
 
   it('should work with a dispatched state change for a useReducer', () => {
-    const initialState = {count: 0};
-
     function reducer(state, action) {
       switch (action.type) {
-        case 'reset':
-          return initialState;
         case 'increment':
           return {count: state.count + 1};
         case 'decrement':
           return {count: state.count - 1};
-        default:
-          return state;
       }
     }
 
-    function SomeComponent({initialCount}) {
-      const [state, dispatch] = React.useReducer(reducer, {
-        count: initialCount,
-      });
+    function SomeComponent(props) {
+      const [state, dispatch] = React.useReducer(reducer, props, p => ({
+        count: p.initialCount,
+      }));
 
       if (state.count === 0) {
         dispatch({type: 'increment'});