From 9907718e09a7f85074f5bfd8109caeb1ebc2c186 Mon Sep 17 00:00:00 2001
From: Jonathan Ziller <jonathan.ziller@gmail.com>
Date: Sat, 11 Jul 2020 11:29:28 +0200
Subject: [PATCH] fix: make `wrapReducerWithFormStateUpdate` work properly if
 used on states that are form states themselves

---
 CHANGELOG.md        |  7 +++++++
 src/reducer.spec.ts | 46 +++++++++++++++++++++++++++++++++++++++++++++
 src/reducer.ts      |  6 ++++++
 3 files changed, 59 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a42b1825..dd70036f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,12 @@
 ## ngrx-forms Changelog
 
+<a name="6.3.3"></a>
+### 6.3.3
+
+#### Bugfixes
+
+* make `wrapReducerWithFormStateUpdate` work properly if used on states that are form states themselves, closes [#196](https://github.com/MrWolfZ/ngrx-forms/issues/196)
+
 <a name="6.3.2"></a>
 ### 6.3.2
 
diff --git a/src/reducer.spec.ts b/src/reducer.spec.ts
index b705714f..0ab0f577 100644
--- a/src/reducer.spec.ts
+++ b/src/reducer.spec.ts
@@ -1,6 +1,10 @@
 import { Action, createReducer } from '@ngrx/store';
 import { MarkAsDirtyAction, MarkAsTouchedAction, SetValueAction } from './actions';
+import { formArrayReducer } from './array/reducer';
+import { formControlReducer } from './control/reducer';
+import { formGroupReducer } from './group/reducer';
 import { createFormStateReducerWithUpdate, formStateReducer, onNgrxForms, onNgrxFormsAction, wrapReducerWithFormStateUpdate } from './reducer';
+import { FormArrayState, FormControlState, FormGroupState } from './state';
 import { FORM_CONTROL_ID, FORM_CONTROL_INNER5_ID, FORM_CONTROL_INNER_ID, FormGroupValue, INITIAL_STATE } from './update-function/test-util';
 import { updateGroup } from './update-function/update-group';
 
@@ -294,6 +298,20 @@ describe(wrapReducerWithFormStateUpdate.name, () => {
     expect(resultState.control).not.toBe(INITIAL_STATE.controls.inner);
   });
 
+  it('should update a non-nested control after the reducer', () => {
+    const wrappedReducer = wrapReducerWithFormStateUpdate<FormControlState<string>, FormControlState<string>>(
+      formControlReducer,
+      s => s,
+      s => {
+        expect(s).toBe(INITIAL_STATE.controls.inner);
+        return ({ ...s });
+      },
+    );
+
+    const resultState = wrappedReducer(initialState.control, { type: '' });
+    expect(resultState).not.toBe(INITIAL_STATE.controls.inner);
+  });
+
   it('should update a group after the reducer', () => {
     const wrappedReducer = wrapReducerWithFormStateUpdate(reducer, s => s.group, s => {
       expect(s).toBe(INITIAL_STATE);
@@ -304,6 +322,20 @@ describe(wrapReducerWithFormStateUpdate.name, () => {
     expect(resultState.group).not.toBe(INITIAL_STATE);
   });
 
+  it('should update a non-nested group after the reducer', () => {
+    const wrappedReducer = wrapReducerWithFormStateUpdate<FormGroupState<FormGroupValue>, FormGroupState<FormGroupValue>>(
+      formGroupReducer,
+      s => s,
+      s => {
+        expect(s).toBe(INITIAL_STATE);
+        return ({ ...s });
+      },
+    );
+
+    const resultState = wrappedReducer(initialState.group, { type: '' });
+    expect(resultState).not.toBe(INITIAL_STATE);
+  });
+
   it('should update an array after the reducer', () => {
     const wrappedReducer = wrapReducerWithFormStateUpdate(reducer, s => s.array, s => {
       expect(s).toBe(INITIAL_STATE.controls.inner5);
@@ -314,6 +346,20 @@ describe(wrapReducerWithFormStateUpdate.name, () => {
     expect(resultState.array).not.toBe(INITIAL_STATE.controls.inner5);
   });
 
+  it('should update a non-nested array after the reducer', () => {
+    const wrappedReducer = wrapReducerWithFormStateUpdate<FormArrayState<string>, FormArrayState<string>>(
+      formArrayReducer,
+      s => s,
+      s => {
+        expect(s).toBe(INITIAL_STATE.controls.inner5);
+        return ({ ...s });
+      },
+    );
+
+    const resultState = wrappedReducer(initialState.array, { type: '' });
+    expect(resultState).not.toBe(INITIAL_STATE.controls.inner5);
+  });
+
   it('should set the updated form state', () => {
     const updatedControl = { ...INITIAL_STATE.controls.inner };
     const wrappedReducer = wrapReducerWithFormStateUpdate(reducer, s => s.control, () => updatedControl);
diff --git a/src/reducer.ts b/src/reducer.ts
index 526e8ba0..689bd656 100644
--- a/src/reducer.ts
+++ b/src/reducer.ts
@@ -150,6 +150,12 @@ export function wrapReducerWithFormStateUpdate<TState, TFormState extends Abstra
     const updatedState = reducer(state, action);
 
     const formState = formStateLocator(updatedState);
+
+    // if the state itself is the form state, update it directly
+    if (formState === updatedState as unknown) {
+      return updateFn(formState, updatedState) as unknown as TState;
+    }
+
     const formStateKey = Object.keys(updatedState).find(key => updatedState[key as keyof TState] as any === formState)!;
 
     const updatedFormState = updateFn(formState, updatedState);