From fd14a7d6eb454a3d5035669c7d94b9709586dff7 Mon Sep 17 00:00:00 2001
From: Vladimir Gorej <vladimir.gorej@gmail.com>
Date: Mon, 24 May 2021 14:49:18 +0200
Subject: [PATCH] fix(system): allow wrapping selectors in chain

Refs #7157
---
 src/core/system.js                     | 47 +++++++++++++++++++-------
 test/unit/core/system/wrapSelectors.js | 45 ++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 12 deletions(-)
 create mode 100644 test/unit/core/system/wrapSelectors.js

diff --git a/src/core/system.js b/src/core/system.js
index a756039c975..3179b43fd71 100644
--- a/src/core/system.js
+++ b/src/core/system.js
@@ -392,23 +392,46 @@ function systemExtend(dest={}, src={}) {
   if(isObject(statePlugins)) {
     for(let namespace in statePlugins) {
       const namespaceObj = statePlugins[namespace]
-      if(!isObject(namespaceObj) || !isObject(namespaceObj.wrapActions)) {
+      if(!isObject(namespaceObj)) {
         continue
       }
-      const { wrapActions } = namespaceObj
-      for(let actionName in wrapActions) {
-        let action = wrapActions[actionName]
-
-        // This should only happen if dest is the first plugin, since invocations after that will ensure its an array
-        if(!Array.isArray(action)) {
-          action = [action]
-          wrapActions[actionName] = action // Put the value inside an array
-        }
 
-        if(src && src.statePlugins && src.statePlugins[namespace] && src.statePlugins[namespace].wrapActions && src.statePlugins[namespace].wrapActions[actionName]) {
-          src.statePlugins[namespace].wrapActions[actionName] = wrapActions[actionName].concat(src.statePlugins[namespace].wrapActions[actionName])
+      const { wrapActions, wrapSelectors } = namespaceObj
+
+      // process action wrapping
+      if (isObject(wrapActions)) {
+        for(let actionName in wrapActions) {
+          let action = wrapActions[actionName]
+
+          // This should only happen if dest is the first plugin, since invocations after that will ensure its an array
+          if(!Array.isArray(action)) {
+            action = [action]
+            wrapActions[actionName] = action // Put the value inside an array
+          }
+
+          if(src && src.statePlugins && src.statePlugins[namespace] && src.statePlugins[namespace].wrapActions && src.statePlugins[namespace].wrapActions[actionName]) {
+            src.statePlugins[namespace].wrapActions[actionName] = wrapActions[actionName].concat(src.statePlugins[namespace].wrapActions[actionName])
+          }
+
         }
+      }
+
+      // process selector wrapping
+      if (isObject(wrapSelectors)) {
+        for(let selectorName in wrapSelectors) {
+          let selector = wrapSelectors[selectorName]
+
+          // This should only happen if dest is the first plugin, since invocations after that will ensure its an array
+          if(!Array.isArray(selector)) {
+            selector = [selector]
+            wrapSelectors[selectorName] = selector // Put the value inside an array
+          }
 
+          if(src && src.statePlugins && src.statePlugins[namespace] && src.statePlugins[namespace].wrapSelectors && src.statePlugins[namespace].wrapSelectors[selectorName]) {
+            src.statePlugins[namespace].wrapSelectors[selectorName] = wrapSelectors[selectorName].concat(src.statePlugins[namespace].wrapSelectors[selectorName])
+          }
+
+        }
       }
     }
   }
diff --git a/test/unit/core/system/wrapSelectors.js b/test/unit/core/system/wrapSelectors.js
new file mode 100644
index 00000000000..076dc0dcc02
--- /dev/null
+++ b/test/unit/core/system/wrapSelectors.js
@@ -0,0 +1,45 @@
+import System from "core/system"
+
+describe("wrapSelectors", () => {
+  it("should wrap correctly when registering multiple plugins targeting the same selector", function() {
+    const probeBase = {
+      statePlugins: {
+        probe: {
+          selectors: {
+            selectProbe: () => {
+              return "base"
+            }
+          }
+        }
+      }
+    }
+    const probeWrap1 = {
+      statePlugins: {
+        probe: {
+          wrapSelectors: {
+            selectProbe: (oriSelector) => (state, ...args) => {
+              const selectedValue = oriSelector(state, ...args)
+              return `${selectedValue}wrap1`
+            }
+          }
+        }
+      }
+    }
+    const probeWrap2 = {
+      statePlugins: {
+        probe: {
+          wrapSelectors: {
+            selectProbe: (oriSelector) => (state, ...args) => {
+              const selectedValue = oriSelector(state, ...args)
+              return `${selectedValue}wrap2`
+            }
+          }
+        }
+      }
+    }
+
+    const system = new System({ plugins: [probeBase, probeWrap1, probeWrap2] })
+
+    expect(system.getSystem().probeSelectors.selectProbe()).toEqual("basewrap1wrap2")
+  })
+})