From e6c28e0d459ba543b747e315a9a4c2d539c24de6 Mon Sep 17 00:00:00 2001
From: Yoshua Wuyts <yoshuawuyts@gmail.com>
Date: Wed, 7 Sep 2016 17:21:29 +0200
Subject: [PATCH] fixup! fix statewrap

---
 index.js | 70 +++++++++++++++++++++++++++++++++++++-------------------
 test.js  | 13 +++++++----
 2 files changed, 55 insertions(+), 28 deletions(-)

diff --git a/index.js b/index.js
index 6eb6477..e311e82 100644
--- a/index.js
+++ b/index.js
@@ -52,7 +52,7 @@ function dispatcher (hooks) {
     if (hooks.onError) onErrorHooks.push(wrapOnError(hooks.onError))
     if (hooks.onAction) onActionHooks.push(hooks.onAction)
     if (hooks.wrapSubscriptions) subscriptionWraps.push(hooks.wrapSubscriptions)
-    if (hooks.initialState) initialStateWraps.push(hooks.wrapInitialState)
+    if (hooks.wrapInitialState) initialStateWraps.push(hooks.wrapInitialState)
     if (hooks.wrapReducers) reducerWraps.push(hooks.wrapReducers)
     if (hooks.wrapEffects) effectWraps.push(hooks.wrapEffects)
   }
@@ -69,25 +69,30 @@ function dispatcher (hooks) {
   function getState (opts) {
     opts = opts || {}
     assert.equal(typeof opts, 'object', 'barracks.store.state: opts should be an object')
-    if (opts.state) {
-      const initialState = {}
-      const nsState = {}
-      models.forEach(function (model) {
-        const ns = model.namespace
-        const modelState = model.state || {}
-        if (ns) {
-          nsState[ns] = {}
-          apply(ns, modelState, nsState)
-          nsState[ns] = xtend(nsState[ns], opts.state[ns])
-        } else {
-          apply(model.namespace, modelState, initialState)
-        }
-      })
-      return xtend(_state, xtend(opts.state, nsState))
-    } else if (opts.freeze === false) {
-      return xtend(_state)
+
+    const state = opts.state
+    if (!opts.state && opts.freeze === false) return xtend(_state)
+    else if (!opts.state) return Object.freeze(xtend(_state))
+    assert.equal(typeof state, 'object', 'barracks.store.state: state should be an object')
+
+    const nsState = {}
+
+    models.forEach(function (model) {
+      const ns = model.namespace
+      const modelState = model.state || {}
+      if (ns) {
+        nsState[ns] = {}
+        apply(ns, modelState, nsState)
+        nsState[ns] = xtend(nsState[ns], state[ns])
+      } else {
+        mutate(nsState, modelState)
+      }
+    })
+
+    if (opts.freeze === false) {
+      return xtend(_state, xtend(state, nsState))
     } else {
-      return Object.freeze(xtend(_state))
+      return Object.freeze(xtend(_state, xtend(state, nsState)))
     }
   }
 
@@ -100,6 +105,15 @@ function dispatcher (hooks) {
     // register values from the models
     models.forEach(function (model) {
       const ns = model.namespace
+      if (!stateCalled && model.state && opts.state !== false) {
+        const modelState = model.state || {}
+        if (ns) {
+          _state[ns] = _state[ns] || {}
+          apply(ns, modelState, _state)
+        } else {
+          mutate(_state, modelState)
+        }
+      }
       if (!reducersCalled && model.reducers && opts.reducers !== false) {
         apply(ns, model.reducers, reducers, function (cb) {
           return wrapHook(cb, reducerWraps)
@@ -122,9 +136,16 @@ function dispatcher (hooks) {
       }
     })
 
+    // the state wrap is special because we want to operate on the full
+    // state rather than indvidual chunks, so we apply it outside the loop
+    if (!stateCalled && opts.state !== false) {
+      _state = wrapHook(_state, initialStateWraps)
+    }
+
     if (!opts.noSubscriptions) subsCalled = true
     if (!opts.noReducers) reducersCalled = true
     if (!opts.noEffects) effectsCalled = true
+    if (!opts.noState) stateCalled = true
 
     if (!onErrorHooks.length) onErrorHooks.push(wrapOnError(defaultOnError))
 
@@ -245,9 +266,12 @@ function wrapOnError (onError) {
   }
 }
 
-function wrapHook (fn, arr) {
-  arr.forEach(function (cb) {
-    fn = cb(fn)
+// take a apply an array of transforms onto a value. The new value
+// must be returned synchronously from the transform
+// (any, [fn]) -> any
+function wrapHook (value, transforms) {
+  transforms.forEach(function (transform) {
+    value = transform(value)
   })
-  return fn
+  return value
 }
diff --git a/test.js b/test.js
index 2b49cbd..40311f8 100644
--- a/test.js
+++ b/test.js
@@ -1,4 +1,5 @@
 const barracks = require('./')
+const xtend = require('xtend')
 const noop = require('noop2')
 const tape = require('tape')
 
@@ -554,12 +555,12 @@ tape('wrappers: wrapEffects')
 
 tape('wrappers: wrapInitialState', (t) => {
   t.test('should wrap initial state in start', (t) => {
-    t.plan(1)
+    t.plan(2)
     const store = barracks()
     store.use({
       wrapInitialState: (state) => {
-        console.log('oi')
-        console.log('state', state)
+        t.deepEqual(state, { foo: 'bar' }, 'initial state is correct')
+        return xtend(state, { beep: 'boop' })
       }
     })
 
@@ -568,8 +569,10 @@ tape('wrappers: wrapInitialState', (t) => {
     })
 
     store.start()
-    const state = store.state()
-    console.log('yo', state)
+    process.nextTick(() => {
+      const state = store.state()
+      t.deepEqual(state, { foo: 'bar', beep: 'boop' }, 'wrapped state correct')
+    })
   })
 
   t.test('should wrap initial state in getState', (t) => {