From 291379c27278cf1bde1a16e2184e750cd1a15a92 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Mon, 28 Mar 2016 21:05:05 +0100
Subject: [PATCH] Elements from functional components in TestUtils should have
 no owner

Fixes #5292 and #6194
---
 .../reconciler/ReactCompositeComponent.js     | 76 ++++++++++---------
 src/test/ReactTestUtils.js                    |  2 +
 src/test/__tests__/ReactTestUtils-test.js     | 20 +++++
 3 files changed, 62 insertions(+), 36 deletions(-)

diff --git a/src/renderers/shared/reconciler/ReactCompositeComponent.js b/src/renderers/shared/reconciler/ReactCompositeComponent.js
index 9e636294b55d8..42596e9b2504c 100644
--- a/src/renderers/shared/reconciler/ReactCompositeComponent.js
+++ b/src/renderers/shared/reconciler/ReactCompositeComponent.js
@@ -61,6 +61,10 @@ function warnIfInvalidElement(Component, element) {
   }
 }
 
+function shouldConstruct(Component) {
+  return Component.prototype && Component.prototype.isReactComponent;
+}
+
 /**
  * ------------------ The Life-Cycle of a Composite Component ------------------
  *
@@ -159,44 +163,22 @@ var ReactCompositeComponentMixin = {
     var Component = this._currentElement.type;
 
     // Initialize the public class
-    var inst;
+    var inst = this._constructComponent(publicProps, publicContext);
     var renderedElement;
 
-    if (Component.prototype && Component.prototype.isReactComponent) {
-      if (__DEV__) {
-        ReactCurrentOwner.current = this;
-        try {
-          inst = new Component(publicProps, publicContext, ReactUpdateQueue);
-        } finally {
-          ReactCurrentOwner.current = null;
-        }
-      } else {
-        inst = new Component(publicProps, publicContext, ReactUpdateQueue);
-      }
-    } else {
-      if (__DEV__) {
-        ReactCurrentOwner.current = this;
-        try {
-          inst = Component(publicProps, publicContext, ReactUpdateQueue);
-        } finally {
-          ReactCurrentOwner.current = null;
-        }
-      } else {
-        inst = Component(publicProps, publicContext, ReactUpdateQueue);
-      }
-      if (inst == null || inst.render == null) {
-        renderedElement = inst;
-        warnIfInvalidElement(Component, renderedElement);
-        invariant(
-          inst === null ||
-          inst === false ||
-          ReactElement.isValidElement(inst),
-          '%s(...): A valid React element (or null) must be returned. You may have ' +
-          'returned undefined, an array or some other invalid object.',
-          Component.displayName || Component.name || 'Component'
-        );
-        inst = new StatelessComponent(Component);
-      }
+    // Support functional components
+    if (!shouldConstruct(Component) && (inst == null || inst.render == null)) {
+      renderedElement = inst;
+      warnIfInvalidElement(Component, renderedElement);
+      invariant(
+        inst === null ||
+        inst === false ||
+        ReactElement.isValidElement(inst),
+        '%s(...): A valid React element (or null) must be returned. You may have ' +
+        'returned undefined, an array or some other invalid object.',
+        Component.displayName || Component.name || 'Component'
+      );
+      inst = new StatelessComponent(Component);
     }
 
     if (__DEV__) {
@@ -324,6 +306,28 @@ var ReactCompositeComponentMixin = {
     return markup;
   },
 
+  _constructComponent: function(publicProps, publicContext) {
+    if (__DEV__) {
+      ReactCurrentOwner.current = this;
+      try {
+        return this._constructComponentWithoutOwner(publicProps, publicContext);
+      } finally {
+        ReactCurrentOwner.current = null;
+      }
+    } else {
+      return this._constructComponentWithoutOwner(publicProps, publicContext);
+    }
+  },
+
+  _constructComponentWithoutOwner: function(publicProps, publicContext) {
+    var Component = this._currentElement.type;
+    if (shouldConstruct(Component)) {
+      return new Component(publicProps, publicContext, ReactUpdateQueue);
+    } else {
+      return Component(publicProps, publicContext, ReactUpdateQueue);
+    }
+  },
+
   performInitialMountWithErrorHandling: function(
     renderedElement,
     nativeParent,
diff --git a/src/test/ReactTestUtils.js b/src/test/ReactTestUtils.js
index 4f8e4ebb38c2b..187e8402b65a8 100644
--- a/src/test/ReactTestUtils.js
+++ b/src/test/ReactTestUtils.js
@@ -399,6 +399,8 @@ var ShallowComponentWrapper = function(element) {
 assign(
   ShallowComponentWrapper.prototype,
   ReactCompositeComponent.Mixin, {
+    _constructComponent:
+      ReactCompositeComponent.Mixin._constructComponentWithoutOwner,
     _instantiateReactComponent: function(element) {
       return new NoopInternalComponent(element);
     },
diff --git a/src/test/__tests__/ReactTestUtils-test.js b/src/test/__tests__/ReactTestUtils-test.js
index a7cf4ab846536..07cc4a6605ca4 100644
--- a/src/test/__tests__/ReactTestUtils-test.js
+++ b/src/test/__tests__/ReactTestUtils-test.js
@@ -47,6 +47,26 @@ describe('ReactTestUtils', function() {
     ]);
   });
 
+  it('should shallow render a functional component', function() {
+    function SomeComponent() {
+      return (
+        <div>
+          <span className="child1" />
+          <span className="child2" />
+        </div>
+      );
+    }
+
+    var shallowRenderer = ReactTestUtils.createRenderer();
+    var result = shallowRenderer.render(<SomeComponent />);
+
+    expect(result.type).toBe('div');
+    expect(result.props.children).toEqual([
+      <span className="child1" />,
+      <span className="child2" />,
+    ]);
+  });
+
   it('should throw for invalid elements', function() {
     var SomeComponent = React.createClass({
       render: function() {