From 66f66622eb037f4783caa915d0282df486a57936 Mon Sep 17 00:00:00 2001
From: Jordan Harband <ljharb@gmail.com>
Date: Thu, 1 Jun 2017 15:45:34 -0700
Subject: [PATCH] [Refactor] `shallow`/`mount`: Make all references to the
 wrapper `class` call into `this.wrap`

---
 src/ReactWrapper.jsx  | 4 ++--
 src/ShallowWrapper.js | 8 ++++----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/ReactWrapper.jsx b/src/ReactWrapper.jsx
index 95cc1d599..246a4a7c4 100644
--- a/src/ReactWrapper.jsx
+++ b/src/ReactWrapper.jsx
@@ -924,11 +924,11 @@ class ReactWrapper {
    * @param {ReactWrapper|ReactElement|Array<ReactElement>} node
    * @returns {ReactWrapper}
    */
-  wrap(node) {
+  wrap(node, root = this.root, ...args) {
     if (node instanceof ReactWrapper) {
       return node;
     }
-    return new ReactWrapper(node, this.root);
+    return new ReactWrapper(node, root, ...args);
   }
 
   /**
diff --git a/src/ShallowWrapper.js b/src/ShallowWrapper.js
index f0428e5d0..370a1b820 100644
--- a/src/ShallowWrapper.js
+++ b/src/ShallowWrapper.js
@@ -755,7 +755,7 @@ class ShallowWrapper {
    * @returns {ShallowWrapper}
    */
   shallow(options) {
-    return this.single('shallow', n => new ShallowWrapper(n, null, options));
+    return this.single('shallow', n => this.wrap(n, null, options));
   }
 
   /**
@@ -1039,11 +1039,11 @@ class ShallowWrapper {
    * @param node
    * @returns {ShallowWrapper}
    */
-  wrap(node) {
+  wrap(node, root = this.root, ...args) {
     if (node instanceof ShallowWrapper) {
       return node;
     }
-    return new ShallowWrapper(node, this.root);
+    return new ShallowWrapper(node, root, ...args);
   }
 
   /**
@@ -1082,7 +1082,7 @@ class ShallowWrapper {
       if (!isCustomComponentElement(n)) {
         throw new TypeError(`ShallowWrapper::${name}() can only be called on components`);
       }
-      return new ShallowWrapper(n, null, { ...this.options, ...options });
+      return this.wrap(n, null, { ...this.options, ...options });
     });
   }
 }