stub instance.setState when the component is mounting

Author: koba04

packages/enzyme/src/ShallowWrapper.js

@@ -171,6 +171,13 @@ class ShallowWrapper {
       privateSet(this, RENDERER, renderer);
       this[RENDERER].render(nodes, options.context);
       const { instance } = this[RENDERER].getNode();
+      const adapter = getAdapter(this[OPTIONS]);
+      const lifecycles = getAdapterLifecycles(adapter);
+      // Ensure to call componentDidUpdate when instance.setState is called
+      if (instance && lifecycles.componentDidUpdate.onSetState && !instance[SET_STATE]) {
+        privateSet(instance, SET_STATE, instance.setState);
+        instance.setState = (...args) => this.setState(...args);
+      }
       if (
         !options.disableLifecycleMethods
         && instance
@@ -264,16 +271,7 @@ class ShallowWrapper {
     if (this[ROOT] !== this) {
       throw new Error('ShallowWrapper::instance() can only be called on the root');
     }
-    const adapter = getAdapter(this[OPTIONS]);
-    const lifecycles = getAdapterLifecycles(adapter);
-
-    const { instance } = this[RENDERER].getNode();
-    // Ensure to call componentDidUpdate when instance.setState is called
-    if (instance && lifecycles.componentDidUpdate.onSetState && !instance[SET_STATE]) {
-      privateSet(instance, SET_STATE, instance.setState);
-      instance.setState = (...args) => this.setState(...args);
-    }
-    return instance;
+    return this[RENDERER].getNode().instance;
   }
 
   /**