ReactTestUtils.mockComponent() supports function and forwardRef components

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', () => {

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;
   },