From 074f2ed6916165674bf5bd620d7e5262d2d16169 Mon Sep 17 00:00:00 2001
From: Adam Zmenak <adam.zmenak@gmail.com>
Date: Fri, 20 Dec 2019 15:43:56 -0500
Subject: [PATCH] Fixes #1

---
 src/index.spec.tsx | 54 ++++++++++++++++++++++++++++++++++++++++++++++
 src/index.tsx      |  8 +++++++
 2 files changed, 62 insertions(+)

diff --git a/src/index.spec.tsx b/src/index.spec.tsx
index 8ed5977..6325c62 100644
--- a/src/index.spec.tsx
+++ b/src/index.spec.tsx
@@ -103,6 +103,60 @@ describe('useSagaReducer()', () => {
     expect(el.textContent).toBe('2')
   })
 
+  it('saga updates the state available to yield select()', async () => {
+    let testState: any
+    const testReducer = jest.fn((state = {}, action: any) => {
+      if (action.type === 'INCREMENT') {
+        return {
+          count: state.count + 1
+        }
+      }
+
+      return state
+    })
+
+    function* increment() {
+      const state = yield select()
+      testState = state
+    }
+
+    function* testSaga() {
+      yield takeEvery('INCREMENT', increment)
+    }
+
+    function TestUseSagaReducer() {
+      const [, dispatch] = useSagaReducer(testSaga, testReducer, {count: 1})
+
+      return (
+        <div>
+          <button
+            data-testid='button'
+            onClick={() => {
+              dispatch({
+                type: 'INCREMENT'
+              })
+            }}
+          >
+            TEST
+          </button>
+        </div>
+      )
+    }
+
+    const {getByTestId} = render(<TestUseSagaReducer />)
+    const button = getByTestId('button')
+
+    fireEvent.click(button)
+    await act(flushPromiseQueue)
+
+    expect(testState).toEqual({count: 2})
+
+    fireEvent.click(button)
+    await act(flushPromiseQueue)
+
+    expect(testState).toEqual({count: 3})
+  })
+
   it('provides context values in sagas passed to provider', async () => {
     const testReducer = jest.fn((state = {}, action: any) => {
       return state
diff --git a/src/index.tsx b/src/index.tsx
index 421c3df..f967169 100644
--- a/src/index.tsx
+++ b/src/index.tsx
@@ -4,6 +4,7 @@ import React, {
   useMemo,
   useEffect,
   useContext,
+  useLayoutEffect,
   Reducer,
   ReducerState,
   ReducerAction,
@@ -28,6 +29,9 @@ export const SagaProvider: React.FC<SagaProdiderProps> = (props) => {
   return <SagaContext.Provider {...props} />
 }
 
+const useIsomorphicLayoutEffect =
+  typeof window !== 'undefined' ? useLayoutEffect : useEffect
+
 export function useSagaReducer<
   S extends Saga<never[]>,
   R extends Reducer<any, any>,
@@ -46,6 +50,10 @@ export function useSagaReducer<
   )
 
   const stateRef = useRef(state)
+  useIsomorphicLayoutEffect(() => {
+    stateRef.current = state
+  }, [state])
+
   const sagaIO: Required<Pick<
     RunSagaOptions<any, S>,
     SagaIOKeys