freeze ROOT_NODES

enzymejs · Sep 19, 2018 · 2a4f216 · 2a4f216
1 parent 4dd2780
commit 2a4f216
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/packages/enzyme/src/ReactWrapper.js b/packages/enzyme/src/ReactWrapper.js
@@ -31,6 +31,7 @@ const RENDERER = sym('__renderer__');
 const UNRENDERED = sym('__unrendered__');
 const ROOT = sym('__root__');
 const OPTIONS = sym('__options__');
+const ROOT_NODES = sym('__rootNodes__');
 
 /**
  * Finds all nodes in the current wrapper nodes' render trees that match the provided predicate
@@ -57,8 +58,18 @@ function filterWhereUnwrapped(wrapper, predicate) {
   return wrapper.wrap(wrapper.getNodesInternal().filter(predicate).filter(Boolean));
 }
 
+function getRootNodeInternal(wrapper) {
+  if (wrapper[ROOT].length !== 1) {
+    throw new Error('getRootNodeInternal(wrapper) can only be called when wrapper wraps one node');
+  }
+  if (wrapper[ROOT] !== wrapper) {
+    return wrapper[ROOT_NODES][0];
+  }
+  return wrapper[ROOT][NODE];
+}
+
 function nodeParents(wrapper, node) {
-  return parentsOfNode(node, wrapper[ROOT].getNodeInternal());
+  return parentsOfNode(node, getRootNodeInternal(wrapper));
 }
 
 function privateSetNodes(wrapper, nodes) {
@@ -102,6 +113,7 @@ class ReactWrapper {
       privateSet(this, RENDERER, root[RENDERER]);
       privateSet(this, ROOT, root);
       privateSetNodes(this, nodes);
+      privateSet(this, ROOT_NODES, root[NODES]);
     }
     privateSet(this, OPTIONS, root ? root[OPTIONS] : options);
   }

diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
@@ -37,6 +37,8 @@ const UNRENDERED = sym('__unrendered__');
 const ROOT = sym('__root__');
 const OPTIONS = sym('__options__');
 const SET_STATE = sym('__setState__');
+const ROOT_NODES = sym('__rootNodes__');
+
 /**
  * Finds all nodes in the current wrapper nodes' render trees that match the provided predicate
  * function.
@@ -144,6 +146,12 @@ function getRootNode(node) {
 }
 
 function getRootNodeInternal(wrapper) {
+  if (wrapper[ROOT].length !== 1) {
+    throw new Error('getRootNodeInternal(wrapper) can only be called when wrapper wraps one node');
+  }
+  if (wrapper[ROOT] !== wrapper) {
+    return wrapper[ROOT_NODES][0];
+  }
   return wrapper[ROOT][NODE];
 }
 
@@ -219,6 +227,7 @@ class ShallowWrapper {
       privateSet(this, RENDERER, root[RENDERER]);
       privateSetNodes(this, nodes);
       privateSet(this, OPTIONS, root[OPTIONS]);
+      privateSet(this, ROOT_NODES, root[NODES]);
     }
   }