From 3fcb66621492619bcde6fbbaa3f63e5dcd6ce6b9 Mon Sep 17 00:00:00 2001
From: Brian Vaughn <bvaughn@fb.com>
Date: Wed, 11 Jul 2018 07:42:56 -0700
Subject: [PATCH] ReactTestUtils.mockComponent() support for function and
 forwardRef components

---
 .../src/__tests__/ReactTestUtils-test.js      | 63 +++++++++++++++----
 .../src/test-utils/ReactTestUtils.js          | 20 +++++-
 2 files changed, 68 insertions(+), 15 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactTestUtils-test.js b/packages/react-dom/src/__tests__/ReactTestUtils-test.js
index 83437133b6a6f..ee31b51170d7b 100644
--- a/packages/react-dom/src/__tests__/ReactTestUtils-test.js
+++ b/packages/react-dom/src/__tests__/ReactTestUtils-test.js
@@ -43,21 +43,60 @@ describe('ReactTestUtils', () => {
     expect(Object.keys(ReactTestUtils.SimulateNative).sort()).toMatchSnapshot();
   });
 
-  it('gives Jest mocks a passthrough implementation with mockComponent()', () => {
-    class MockedComponent extends React.Component {
-      render() {
-        throw new Error('Should not get here.');
+  describe('mockComponent', () => {
+    it('should support jest-mocked class components', () => {
+      class MockedClassComponent extends React.Component {
+        render() {
+          throw new Error('Should not get here.');
+        }
       }
-    }
-    // This is close enough to what a Jest mock would give us.
-    MockedComponent.prototype.render = jest.fn();
 
-    // Patch it up so it returns its children.
-    ReactTestUtils.mockComponent(MockedComponent);
+      // This is close enough to what a Jest mock would give us.
+      MockedClassComponent.prototype.render = jest.fn();
 
-    const container = document.createElement('div');
-    ReactDOM.render(<MockedComponent>Hello</MockedComponent>, container);
-    expect(container.textContent).toBe('Hello');
+      ReactTestUtils.mockComponent(MockedClassComponent);
+
+      const container = document.createElement('div');
+      ReactDOM.render(
+        <MockedClassComponent>Hello</MockedClassComponent>,
+        container,
+      );
+      expect(container.textContent).toBe('Hello');
+    });
+
+    it('should support jest-mocked functional components', () => {
+      const MockedFunctionComponent = jest.fn(props => {
+        throw new Error('Should not get here.');
+      });
+
+      ReactTestUtils.mockComponent(MockedFunctionComponent);
+
+      const container = document.createElement('div');
+      ReactDOM.render(
+        <MockedFunctionComponent>Hello</MockedFunctionComponent>,
+        container,
+      );
+      expect(container.textContent).toBe('Hello');
+    });
+
+    it('should support jest-mocked forwardRef object', () => {
+      const RefForwardingFunction = (props, ref) => {
+        throw new Error('Should not get here.');
+      };
+      const RefForwardingComponent = React.forwardRef(RefForwardingFunction);
+
+      // This is close enough to what a Jest mock would give us.
+      RefForwardingComponent.render = jest.fn();
+
+      ReactTestUtils.mockComponent(RefForwardingComponent);
+
+      const container = document.createElement('div');
+      ReactDOM.render(
+        <RefForwardingComponent>Hello</RefForwardingComponent>,
+        container,
+      );
+      expect(container.textContent).toBe('Hello');
+    });
   });
 
   it('can scryRenderedComponentsWithType', () => {
diff --git a/packages/react-dom/src/test-utils/ReactTestUtils.js b/packages/react-dom/src/test-utils/ReactTestUtils.js
index cc0c5a89f7f57..3cd66d32a0a74 100644
--- a/packages/react-dom/src/test-utils/ReactTestUtils.js
+++ b/packages/react-dom/src/test-utils/ReactTestUtils.js
@@ -311,9 +311,23 @@ const ReactTestUtils = {
   mockComponent: function(module, mockTagName) {
     mockTagName = mockTagName || module.mockTagName || 'div';
 
-    module.prototype.render.mockImplementation(function() {
-      return React.createElement(mockTagName, null, this.props.children);
-    });
+    if (module.prototype != null) {
+      // Class or functional component
+      if (typeof module.prototype.render === 'function') {
+        module.prototype.render.mockImplementation(function() {
+          return React.createElement(mockTagName, null, this.props.children);
+        });
+      } else {
+        module.mockImplementation(props =>
+          React.createElement(mockTagName, null, props.children),
+        );
+      }
+    } else if (typeof module.render === 'function') {
+      // Forward ref object
+      module.render.mockImplementation(props =>
+        React.createElement(mockTagName, null, props.children),
+      );
+    }
 
     return this;
   },