From 1f314ddca2f299167a05aa2712f4faed14443c0c Mon Sep 17 00:00:00 2001
From: Kyle Shockey <shockey@users.noreply.github.com>
Date: Wed, 17 May 2017 18:40:58 -0700
Subject: [PATCH] Add wrapSelectors functionality

---
 src/core/system.js  |  32 +++++++++++-
 test/core/system.js | 116 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 147 insertions(+), 1 deletion(-)

diff --git a/src/core/system.js b/src/core/system.js
index 551b7ca1..89856c54 100644
--- a/src/core/system.js
+++ b/src/core/system.js
@@ -76,7 +76,7 @@ export default class Store {
     this.boundSystem = Object.assign({},
         this.getRootInjects(),
         this.getWrappedAndBoundActions(dispatch),
-        this.getBoundSelectors(getState, this.getSystem),
+        this.getWrappedAndBoundSelectors(getState, this.getSystem),
         this.getStateThunks(getState),
         this.getFn(),
         this.getConfigs()
@@ -176,6 +176,36 @@ export default class Store {
       })
   }
 
+  getWrappedAndBoundSelectors(getState, getSystem) {
+    let selectorGroups = this.getBoundSelectors(getState, getSystem)
+      return objMap(selectorGroups, (selectors, selectorGroupName) => {
+        let stateName = [selectorGroupName.slice(0, -9)] // selectors = 9 chars
+        let wrappers = this.system.statePlugins[stateName].wrapSelectors
+          if(wrappers) {
+            return objMap(selectors, (selector, selectorName) => {
+              let wrap = wrappers[selectorName]
+              if(!wrap) {
+                return selector
+              }
+
+              if(!Array.isArray(wrap)) {
+                wrap = [wrap]
+              }
+              return wrap.reduce((acc, fn) => {
+                let wrappedSelector = (...args) => {
+                  return fn(acc, this.getSystem())(getState().getIn(stateName), ...args)
+                }
+                if(!isFn(wrappedSelector)) {
+                  throw new TypeError("wrapSelector needs to return a function that returns a new function (ie the wrapped action)")
+                }
+                return wrappedSelector
+              }, selector || Function.prototype)
+            })
+          }
+        return selectors
+      })
+  }
+
   getStates(state) {
     return Object.keys(this.system.statePlugins).reduce((obj, key) => {
       obj[key] = state.get(key)
diff --git a/test/core/system.js b/test/core/system.js
index 9f546ce0..4736e3a4 100644
--- a/test/core/system.js
+++ b/test/core/system.js
@@ -326,6 +326,122 @@ describe("bound system", function(){
 
     })
 
+    describe('wrapSelectors', () => {
+      it("should wrap a selector and provide a reference to the original", function(){
+
+        // Given
+        const system = new System({
+          plugins: [
+            {
+              statePlugins: {
+                doge: {
+                  selectors: {
+                    wow: () => (system) => {
+                      return "original"
+                    }
+                  }
+                }
+              }
+            },
+            {
+              statePlugins: {
+                doge: {
+                  wrapSelectors: {
+                    wow: (ori) => (system) => {
+                      // Then
+                      return ori() + " wrapper"
+                    }
+                  }
+                }
+              }
+            }
+          ]
+        })
+
+        // When
+        var res = system.getSystem().dogeSelectors.wow(1)
+        expect(res).toEqual("original wrapper")
+
+      })
+
+      it("should provide a live reference to the system to a wrapper", function(done){
+
+        // Given
+        const mySystem = new System({
+          plugins: [
+            {
+              statePlugins: {
+                doge: {
+                  selectors: {
+                    wow: () => (system) => {
+                      return "original"
+                    }
+                  }
+                }
+              }
+            },
+            {
+              statePlugins: {
+                doge: {
+                  wrapSelectors: {
+                    wow: (ori, system) => () => {
+                      // Then
+                      expect(mySystem.getSystem()).toEqual(system.getSystem())
+                      done()
+                      return ori() + " wrapper"
+                    }
+                  }
+                }
+              }
+            }
+          ]
+        })
+
+        mySystem.getSystem().dogeSelectors.wow(1)
+      })
+
+      it("should provide the state as the first argument to the inner function", function(done){
+
+        // Given
+        const mySystem = new System({
+          state: {
+            doge: {
+              abc: "123"
+            }
+          },
+          plugins: [
+            {
+              statePlugins: {
+                doge: {
+                  selectors: {
+                    wow: () => (system) => {
+                      return "original"
+                    }
+                  }
+                }
+              }
+            },
+            {
+              statePlugins: {
+                doge: {
+                  wrapSelectors: {
+                    wow: (ori, system) => (dogeState) => {
+                      // Then
+                      expect(dogeState.toJS().abc).toEqual('123')
+                      done()
+                      return ori() + " wrapper"
+                    }
+                  }
+                }
+              }
+            }
+          ]
+        })
+
+        mySystem.getSystem().dogeSelectors.wow(1)
+      })
+    })
+
   })
 
 })