From 86bfb522191ed24ca0e13c197fe3596292cb1ef3 Mon Sep 17 00:00:00 2001
From: Brian Vaughn <bvaughn@fb.com>
Date: Mon, 17 Apr 2017 11:37:19 -0700
Subject: [PATCH] New ReactShallowRenderer implementation

Replaced shallow renderer based on test renderer with a custom, isolated implementation that's stack and fiber compatible
---
 scripts/fiber/tests-passing.txt               |   2 +
 .../classic/element/ReactElement.js           |   7 -
 .../dom/test/__tests__/ReactTestUtils-test.js | 522 +++++++++---------
 src/renderers/testing/ReactShallowRenderer.js | 217 +++++---
 4 files changed, 420 insertions(+), 328 deletions(-)

diff --git a/scripts/fiber/tests-passing.txt b/scripts/fiber/tests-passing.txt
index be466e2c6d371..f70a16a027cba 100644
--- a/scripts/fiber/tests-passing.txt
+++ b/scripts/fiber/tests-passing.txt
@@ -1560,6 +1560,7 @@ src/renderers/dom/shared/wrappers/__tests__/ReactDOMTextarea-test.js
 * should warn if value and defaultValue are specified
 
 src/renderers/dom/test/__tests__/ReactTestUtils-test.js
+* should only render 1 level deep
 * should have shallow rendering
 * should shallow render a functional component
 * should throw for invalid elements
@@ -1573,6 +1574,7 @@ src/renderers/dom/test/__tests__/ReactTestUtils-test.js
 * can setState in componentWillMount when shallow rendering
 * can pass context when shallowly rendering
 * can fail context when shallowly rendering
+* should warn about propTypes (but only once)
 * can scryRenderedDOMComponentsWithClass with TextComponent
 * can scryRenderedDOMComponentsWithClass with className contains \n
 * can scryRenderedDOMComponentsWithClass with multiple classes
diff --git a/src/isomorphic/classic/element/ReactElement.js b/src/isomorphic/classic/element/ReactElement.js
index cecaf56b4a564..b47d5e247f545 100644
--- a/src/isomorphic/classic/element/ReactElement.js
+++ b/src/isomorphic/classic/element/ReactElement.js
@@ -162,13 +162,6 @@ var ReactElement = function(type, key, ref, self, source, owner, props) {
         writable: false,
         value: source,
       });
-      // _owner attribute would break shallow renderer equality checks.
-      Object.defineProperty(element, '_owner', {
-        configurable: false,
-        enumerable: false,
-        writable: false,
-        value: owner,
-      });
     } else {
       element._store.validated = false;
       element._self = self;
diff --git a/src/renderers/dom/test/__tests__/ReactTestUtils-test.js b/src/renderers/dom/test/__tests__/ReactTestUtils-test.js
index b68200c55a13e..23c3f41970e9b 100644
--- a/src/renderers/dom/test/__tests__/ReactTestUtils-test.js
+++ b/src/renderers/dom/test/__tests__/ReactTestUtils-test.js
@@ -11,8 +11,6 @@
 
 'use strict';
 
-var ReactDOMFeatureFlags = require('ReactDOMFeatureFlags');
-
 var createRenderer;
 var PropTypes;
 var React;
@@ -30,44 +28,21 @@ describe('ReactTestUtils', () => {
     ReactTestUtils = require('ReactTestUtils');
   });
 
-  // Shallow renderer only implemented for Fiber in 16+
-  if (ReactDOMFeatureFlags.useFiber) {
-    it('should only render 1 level deep', () => {
-      function Parent() {
-        return <div><Child/></div>
-      }
-      function Child() {
-        throw Error('This component should not render')
-      }
-
-      var shallowRenderer = createRenderer();
-      shallowRenderer.render(React.createElement(Parent));
-    });
-
-    it('should have shallow rendering', () => {
-      class SomeComponent extends React.Component {
-        render() {
-          return (
-            <div>
-              <span className="child1" />
-              <span className="child2" />
-            </div>
-          );
-        }
-      }
-
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(<SomeComponent />);
+  it('should only render 1 level deep', () => {
+    function Parent() {
+      return <div><Child /></div>;
+    }
+    function Child() {
+      throw Error('This component should not render');
+    }
 
-      expect(result.type).toBe('div');
-      expect(result.props.children).toEqual([
-        <span className="child1" />,
-        <span className="child2" />,
-      ]);
-    });
+    var shallowRenderer = createRenderer();
+    shallowRenderer.render(React.createElement(Parent));
+  });
 
-    it('should shallow render a functional component', () => {
-      function SomeComponent() {
+  it('should have shallow rendering', () => {
+    class SomeComponent extends React.Component {
+      render() {
         return (
           <div>
             <span className="child1" />
@@ -75,249 +50,295 @@ describe('ReactTestUtils', () => {
           </div>
         );
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(<SomeComponent />);
+    var shallowRenderer = createRenderer();
+    var result = shallowRenderer.render(<SomeComponent />);
 
-      expect(result.type).toBe('div');
-      expect(result.props.children).toEqual([
-        <span className="child1" />,
-        <span className="child2" />,
-      ]);
-    });
+    expect(result.type).toBe('div');
+    expect(result.props.children).toEqual([
+      <span className="child1" />,
+      <span className="child2" />,
+    ]);
+  });
 
-    it('should throw for invalid elements', () => {
-      class SomeComponent extends React.Component {
-        render() {
-          return <div />;
-        }
+  it('should shallow render a functional component', () => {
+    function SomeComponent() {
+      return (
+        <div>
+          <span className="child1" />
+          <span className="child2" />
+        </div>
+      );
+    }
+
+    var shallowRenderer = 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', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return <div />;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      expect(() => shallowRenderer.render(SomeComponent)).toThrowError(
-        'ReactShallowRenderer render(): Invalid component element. Instead of ' +
-          'passing a component class, make sure to instantiate it by passing it ' +
-          'to React.createElement.',
-      );
-      expect(() => shallowRenderer.render(<div />)).toThrowError(
-        'ReactShallowRenderer render(): Shallow rendering works only with ' +
-          'custom components, not primitives (div). Instead of calling ' +
-          '`.render(el)` and inspecting the rendered output, look at `el.props` ' +
-          'directly instead.',
-      );
-    });
+    var shallowRenderer = createRenderer();
+    expect(() => shallowRenderer.render(SomeComponent)).toThrowError(
+      'ReactShallowRenderer render(): Invalid component element. Instead of ' +
+        'passing a component class, make sure to instantiate it by passing it ' +
+        'to React.createElement.',
+    );
+    expect(() => shallowRenderer.render(<div />)).toThrowError(
+      'ReactShallowRenderer render(): Shallow rendering works only with ' +
+        'custom components, not primitives (div). Instead of calling ' +
+        '`.render(el)` and inspecting the rendered output, look at `el.props` ' +
+        'directly instead.',
+    );
+  });
 
-    it('should have shallow unmounting', () => {
-      var componentWillUnmount = jest.fn();
+  it('should have shallow unmounting', () => {
+    var componentWillUnmount = jest.fn();
 
-      class SomeComponent extends React.Component {
-        componentWillUnmount = componentWillUnmount;
-        render() {
-          return <div />;
-        }
+    class SomeComponent extends React.Component {
+      componentWillUnmount = componentWillUnmount;
+      render() {
+        return <div />;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      shallowRenderer.render(<SomeComponent />);
-      shallowRenderer.unmount();
+    var shallowRenderer = createRenderer();
+    shallowRenderer.render(<SomeComponent />);
+    shallowRenderer.unmount();
 
-      expect(componentWillUnmount).toBeCalled();
-    });
+    expect(componentWillUnmount).toBeCalled();
+  });
 
-    it('can shallow render to null', () => {
-      class SomeComponent extends React.Component {
-        render() {
-          return null;
-        }
+  it('can shallow render to null', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return null;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(<SomeComponent />);
+    var shallowRenderer = createRenderer();
+    var result = shallowRenderer.render(<SomeComponent />);
 
-      expect(result).toBe(null);
-    });
+    expect(result).toBe(null);
+  });
 
-    it('can shallow render with a ref', () => {
-      class SomeComponent extends React.Component {
-        render() {
-          return <div ref="hello" />;
-        }
+  it('can shallow render with a ref', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return <div ref="hello" />;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      // Shouldn't crash.
-      shallowRenderer.render(<SomeComponent />);
-    });
+    var shallowRenderer = createRenderer();
+    // Shouldn't crash.
+    shallowRenderer.render(<SomeComponent />);
+  });
 
-    it('lets you update shallowly rendered components', () => {
-      class SomeComponent extends React.Component {
-        state = {clicked: false};
-
-        onClick = () => {
-          this.setState({clicked: true});
-        };
-
-        render() {
-          var className = this.state.clicked ? 'was-clicked' : '';
-
-          if (this.props.aNew === 'prop') {
-            return (
-              <a href="#" onClick={this.onClick} className={className}>
-                Test link
-              </a>
-            );
-          } else {
-            return (
-              <div>
-                <span className="child1" />
-                <span className="child2" />
-              </div>
-            );
-          }
+  it('lets you update shallowly rendered components', () => {
+    class SomeComponent extends React.Component {
+      state = {clicked: false};
+
+      onClick = () => {
+        this.setState({clicked: true});
+      };
+
+      render() {
+        var className = this.state.clicked ? 'was-clicked' : '';
+
+        if (this.props.aNew === 'prop') {
+          return (
+            <a href="#" onClick={this.onClick} className={className}>
+              Test link
+            </a>
+          );
+        } else {
+          return (
+            <div>
+              <span className="child1" />
+              <span className="child2" />
+            </div>
+          );
         }
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(<SomeComponent />);
-      expect(result.type).toBe('div');
-      expect(result.props.children).toEqual([
-        <span className="child1" />,
-        <span className="child2" />,
-      ]);
+    var shallowRenderer = createRenderer();
+    var result = shallowRenderer.render(<SomeComponent />);
+    expect(result.type).toBe('div');
+    expect(result.props.children).toEqual([
+      <span className="child1" />,
+      <span className="child2" />,
+    ]);
 
-      var updatedResult = shallowRenderer.render(<SomeComponent aNew="prop" />);
-      expect(updatedResult.type).toBe('a');
+    var updatedResult = shallowRenderer.render(<SomeComponent aNew="prop" />);
+    expect(updatedResult.type).toBe('a');
 
-      var mockEvent = {};
-      updatedResult.props.onClick(mockEvent);
+    var mockEvent = {};
+    updatedResult.props.onClick(mockEvent);
 
-      var updatedResultCausedByClick = shallowRenderer.getRenderOutput();
-      expect(updatedResultCausedByClick.type).toBe('a');
-      expect(updatedResultCausedByClick.props.className).toBe('was-clicked');
-    });
+    var updatedResultCausedByClick = shallowRenderer.getRenderOutput();
+    expect(updatedResultCausedByClick.type).toBe('a');
+    expect(updatedResultCausedByClick.props.className).toBe('was-clicked');
+  });
 
-    it('can access the mounted component instance', () => {
-      class SimpleComponent extends React.Component {
-        someMethod = () => {
-          return this.props.n;
-        };
+  it('can access the mounted component instance', () => {
+    class SimpleComponent extends React.Component {
+      someMethod = () => {
+        return this.props.n;
+      };
 
-        render() {
-          return <div>{this.props.n}</div>;
-        }
+      render() {
+        return <div>{this.props.n}</div>;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      shallowRenderer.render(<SimpleComponent n={5} />);
-      expect(shallowRenderer.getMountedInstance().someMethod()).toEqual(5);
-    });
+    var shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent n={5} />);
+    expect(shallowRenderer.getMountedInstance().someMethod()).toEqual(5);
+  });
 
-    it('can shallowly render components with contextTypes', () => {
-      class SimpleComponent extends React.Component {
-        static contextTypes = {
-          name: PropTypes.string,
-        };
+  it('can shallowly render components with contextTypes', () => {
+    class SimpleComponent extends React.Component {
+      static contextTypes = {
+        name: PropTypes.string,
+      };
 
-        render() {
-          return <div />;
-        }
+      render() {
+        return <div />;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(<SimpleComponent />);
-      expect(result).toEqual(<div />);
-    });
+    var shallowRenderer = createRenderer();
+    var result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div />);
+  });
 
-    it('can shallowly render components with ref as function', () => {
-      class SimpleComponent extends React.Component {
-        state = {clicked: false};
+  it('can shallowly render components with ref as function', () => {
+    class SimpleComponent extends React.Component {
+      state = {clicked: false};
 
-        handleUserClick = () => {
-          this.setState({clicked: true});
-        };
+      handleUserClick = () => {
+        this.setState({clicked: true});
+      };
 
-        render() {
-          return (
-            <div
-              ref={() => {}}
-              onClick={this.handleUserClick}
-              className={this.state.clicked ? 'clicked' : ''}
-            />
-          );
-        }
+      render() {
+        return (
+          <div
+            ref={() => {}}
+            onClick={this.handleUserClick}
+            className={this.state.clicked ? 'clicked' : ''}
+          />
+        );
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      shallowRenderer.render(<SimpleComponent />);
-      var result = shallowRenderer.getRenderOutput();
-      expect(result.type).toEqual('div');
-      expect(result.props.className).toEqual('');
-      result.props.onClick();
+    var shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+    var result = shallowRenderer.getRenderOutput();
+    expect(result.type).toEqual('div');
+    expect(result.props.className).toEqual('');
+    result.props.onClick();
 
-      result = shallowRenderer.getRenderOutput();
-      expect(result.type).toEqual('div');
-      expect(result.props.className).toEqual('clicked');
-    });
+    result = shallowRenderer.getRenderOutput();
+    expect(result.type).toEqual('div');
+    expect(result.props.className).toEqual('clicked');
+  });
 
-    it('can setState in componentWillMount when shallow rendering', () => {
-      class SimpleComponent extends React.Component {
-        componentWillMount() {
-          this.setState({groovy: 'doovy'});
-        }
+  it('can setState in componentWillMount when shallow rendering', () => {
+    class SimpleComponent extends React.Component {
+      componentWillMount() {
+        this.setState({groovy: 'doovy'});
+      }
 
-        render() {
-          return <div>{this.state.groovy}</div>;
-        }
+      render() {
+        return <div>{this.state.groovy}</div>;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(<SimpleComponent />);
-      expect(result).toEqual(<div>doovy</div>);
-    });
+    var shallowRenderer = createRenderer();
+    var result = shallowRenderer.render(<SimpleComponent />);
+    expect(result).toEqual(<div>doovy</div>);
+  });
 
-    it('can pass context when shallowly rendering', () => {
-      class SimpleComponent extends React.Component {
-        static contextTypes = {
-          name: PropTypes.string,
-        };
+  it('can pass context when shallowly rendering', () => {
+    class SimpleComponent extends React.Component {
+      static contextTypes = {
+        name: PropTypes.string,
+      };
 
-        render() {
-          return <div>{this.context.name}</div>;
-        }
+      render() {
+        return <div>{this.context.name}</div>;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(<SimpleComponent />, {
-        name: 'foo',
-      });
-      expect(result).toEqual(<div>foo</div>);
+    var shallowRenderer = createRenderer();
+    var result = shallowRenderer.render(<SimpleComponent />, {
+      name: 'foo',
     });
+    expect(result).toEqual(<div>foo</div>);
+  });
 
-    it('can fail context when shallowly rendering', () => {
-      spyOn(console, 'error');
+  it('can fail context when shallowly rendering', () => {
+    spyOn(console, 'error');
 
-      class SimpleComponent extends React.Component {
-        static contextTypes = {
-          name: PropTypes.string.isRequired,
-        };
+    class SimpleComponent extends React.Component {
+      static contextTypes = {
+        name: PropTypes.string.isRequired,
+      };
 
-        render() {
-          return <div>{this.context.name}</div>;
-        }
+      render() {
+        return <div>{this.context.name}</div>;
       }
+    }
 
-      var shallowRenderer = createRenderer();
-      shallowRenderer.render(<SimpleComponent />);
-      expectDev(console.error.calls.count()).toBe(1);
-      expect(
-        console.error.calls.argsFor(0)[0].replace(/\(at .+?:\d+\)/g, '(at **)'),
-      ).toBe(
-        'Warning: Failed context type: The context `name` is marked as ' +
-          'required in `SimpleComponent`, but its value is `undefined`.\n' +
-          '    in SimpleComponent (at **)',
-      );
-    });
-  } // If fiber
+    var shallowRenderer = createRenderer();
+    shallowRenderer.render(<SimpleComponent />);
+    expectDev(console.error.calls.count()).toBe(1);
+    expect(
+      console.error.calls.argsFor(0)[0].replace(/\(at .+?:\d+\)/g, '(at **)'),
+    ).toBe(
+      'Warning: Failed context type: The context `name` is marked as ' +
+        'required in `SimpleComponent`, but its value is `undefined`.\n' +
+        '    in SimpleComponent (at **)',
+    );
+  });
+
+  it('should warn about propTypes (but only once)', () => {
+    spyOn(console, 'error');
+
+    class SimpleComponent extends React.Component {
+      render() {
+        return React.createElement('div', null, this.props.name);
+      }
+    }
+
+    SimpleComponent.propTypes = {
+      name: PropTypes.string.isRequired,
+    };
+
+    var shallowRenderer = createRenderer();
+    shallowRenderer.render(React.createElement(SimpleComponent, {name: 123}));
+
+    expect(console.error.calls.count()).toBe(1);
+    expect(
+      console.error.calls.argsFor(0)[0].replace(/\(at .+?:\d+\)/g, '(at **)'),
+    ).toBe(
+      'Warning: Failed prop type: Invalid prop `name` of type `number` ' +
+        'supplied to `SimpleComponent`, expected `string`.\n' +
+        '    in SimpleComponent',
+    );
+  });
 
   it('can scryRenderedDOMComponentsWithClass with TextComponent', () => {
     class Wrapper extends React.Component {
@@ -530,32 +551,29 @@ describe('ReactTestUtils', () => {
     );
   });
 
-  // Shallow renderer only implemented for Fiber in 16+
-  if (ReactDOMFeatureFlags.useFiber) {
-    it('should throw when attempting to use ReactTestUtils.Simulate with shallow rendering', () => {
-      class SomeComponent extends React.Component {
-        render() {
-          return (
-            <div onClick={this.props.handleClick}>
-              hello, world.
-            </div>
-          );
-        }
+  it('should throw when attempting to use ReactTestUtils.Simulate with shallow rendering', () => {
+    class SomeComponent extends React.Component {
+      render() {
+        return (
+          <div onClick={this.props.handleClick}>
+            hello, world.
+          </div>
+        );
       }
+    }
 
-      var handler = jasmine.createSpy('spy');
-      var shallowRenderer = createRenderer();
-      var result = shallowRenderer.render(
-        <SomeComponent handleClick={handler} />,
-      );
+    var handler = jasmine.createSpy('spy');
+    var shallowRenderer = createRenderer();
+    var result = shallowRenderer.render(
+      <SomeComponent handleClick={handler} />,
+    );
 
-      expect(() => ReactTestUtils.Simulate.click(result)).toThrowError(
-        'TestUtils.Simulate expects a component instance and not a ReactElement.' +
-          'TestUtils.Simulate will not work if you are using shallow rendering.',
-      );
-      expect(handler).not.toHaveBeenCalled();
-    });
-  }
+    expect(() => ReactTestUtils.Simulate.click(result)).toThrowError(
+      'TestUtils.Simulate expects a component instance and not a ReactElement.' +
+        'TestUtils.Simulate will not work if you are using shallow rendering.',
+    );
+    expect(handler).not.toHaveBeenCalled();
+  });
 
   it('should not warn when simulating events with extra properties', () => {
     spyOn(console, 'error');
diff --git a/src/renderers/testing/ReactShallowRenderer.js b/src/renderers/testing/ReactShallowRenderer.js
index 470b369a8e68e..4df2ec17a075b 100644
--- a/src/renderers/testing/ReactShallowRenderer.js
+++ b/src/renderers/testing/ReactShallowRenderer.js
@@ -12,78 +12,37 @@
 
 'use strict';
 
+const checkPropTypes = require('prop-types/checkPropTypes');
 const React = require('react');
-const ReactTestRenderer = require('react-test-renderer');
 
-const emptyFunction = require('fbjs/lib/emptyFunction');
+const emptyObject = require('fbjs/lib/emptyObject');
 const invariant = require('fbjs/lib/invariant');
 
-const ShallowNodeMockComponent = ({children}) =>
-  React.Children.toArray(children);
+const {
+  ReactDebugCurrentFrame,
+} = React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED;
 
-function createShallowNodeMock() {
-  let isFirst = true;
-  return function createNodeMock(element) {
-    if (isFirst) {
-      isFirst = false;
-      return element.type;
-    }
-    return ShallowNodeMockComponent;
+class ReactShallowRenderer {
+  static createRenderer = function() {
+    return new ReactShallowRenderer();
   };
-}
 
-// TODO Remove this wrapper if/when context parameter is removed
-function wrapElementWithContextProvider(element, context) {
-  const childContextTypes = Object.keys(context).reduce(
-    (reduced, key) => {
-      reduced[key] = emptyFunction;
-      return reduced;
-    },
-    {},
-  );
-
-  class ShallowRendererWrapper extends React.Component {
-    static __shallowRendererWrapperFlag = true;
-    static childContextTypes = childContextTypes;
-    getChildContext() {
-      return context;
-    }
-    render() {
-      return this.props.children;
-    }
+  constructor() {
+    this._element = null;
+    this._instance = null;
+    this._rendered = null;
+    this._updater = new Updater(this);
   }
 
-  return React.createElement(ShallowRendererWrapper, null, element);
-}
-
-class ReactShallowRenderer {
   getMountedInstance() {
-    return this._renderer ? this._renderer.getInstance() : null;
+    return this._instance;
   }
 
   getRenderOutput() {
-    if (this._renderer) {
-      const tree = this._renderer.toTree();
-      if (tree && tree.rendered) {
-        // If we created a context-wrapper then skip over it.
-        const element = tree.type.__shallowRendererWrapperFlag
-          ? tree.rendered.rendered
-          : tree.rendered;
-
-        // Convert the rendered output to a ReactElement.
-        // This supports .toEqual() comparison for test elements.
-        return React.createElement(
-          element.type,
-          element.props,
-          element.props.children,
-        );
-      }
-    }
-    return null;
+    return this._rendered;
   }
 
-  // TODO We should probably remove support for the non-standard context parameter
-  render(element, context) {
+  render(element, context = emptyObject) {
     invariant(
       React.isValidElement(element),
       'ReactShallowRenderer render(): Invalid component element.%s',
@@ -100,26 +59,146 @@ class ReactShallowRenderer {
       element.type,
     );
 
-    // TODO Remove this wrapper if/when context parameter is removed
-    if (context && Object.keys(context).length) {
-      element = wrapElementWithContextProvider(element, context);
-    }
+    this._element = element;
 
-    this._renderer = ReactTestRenderer.create(element, {
-      createNodeMock: createShallowNodeMock(),
-    });
+    if (this._instance) {
+      this._rendered = updateClassComponent(
+        this._instance,
+        element.props,
+        context,
+      );
+    } else {
+      const prototype = element.type.prototype;
+
+      if (typeof prototype.render === 'function') {
+        this._instance = new element.type(
+          element.props,
+          context,
+          this._updater,
+        );
+
+        // TODO context validation: ReactDebugCurrentFrame
+        if (element.type.hasOwnProperty('contextTypes')) {
+          ReactDebugCurrentFrame.element = element;
+
+          checkPropTypes(
+            element.type.contextTypes,
+            context,
+            'context',
+            getName(element.type, this._instance),
+            ReactDebugCurrentFrame.getStackAddendum,
+          );
+
+          ReactDebugCurrentFrame.element = null;
+        }
+
+        this._rendered = mountClassComponent(
+          this._instance,
+          element.props,
+          context,
+        );
+      } else {
+        this._rendered = element.type();
+      }
+    }
 
     return this.getRenderOutput();
   }
 
   unmount() {
-    this._renderer.unmount();
+    if (this._instance) {
+      if (typeof this._instance.componentWillUnmount === 'function') {
+        this._instance.componentWillUnmount();
+      }
+    }
+
+    this._element = null;
+    this._rendered = null;
+    this._instance = null;
   }
 }
 
-// Backwards compatible API
-ReactShallowRenderer.createRenderer = function() {
-  return new ReactShallowRenderer();
-};
+class Updater {
+  constructor(renderer) {
+    this._renderer = renderer;
+  }
+
+  isMounted(publicInstance) {
+    return !!this._renderer._element;
+  }
+
+  enqueueForceUpdate(publicInstance, callback, callerName) {
+    this._renderer.render(this._renderer._element); // TODO
+  }
+
+  enqueueReplaceState(publicInstance, completeState, callback, callerName) {
+    publicInstance.state = completeState;
+    this._renderer.render(this._renderer._element);
+  }
+
+  enqueueSetState(publicInstance, partialState, callback, callerName) {
+    publicInstance.state = {
+      ...publicInstance.state,
+      ...partialState,
+    };
+    this._renderer.render(this._renderer._element);
+  }
+}
+
+function getName(type, instance) {
+  var constructor = instance && instance.constructor;
+  return type.displayName ||
+    (constructor && constructor.displayName) ||
+    type.name ||
+    (constructor && constructor.name) ||
+    null;
+}
+
+function mountClassComponent(
+  instance,
+  props = emptyObject,
+  context = emptyObject,
+) {
+  instance.context = context;
+  instance.props = props;
+  instance.state = instance.state || emptyObject;
+
+  if (typeof instance.componentWillMount === 'function') {
+    instance.componentWillMount();
+  }
+
+  const rendered = instance.render();
+
+  if (typeof instance.componentDidMount === 'function') {
+    instance.componentDidMount();
+  }
+
+  return rendered;
+}
+
+function updateClassComponent(
+  instance,
+  props = emptyObject,
+  context = emptyObject,
+) {
+  const oldContext = instance.context;
+  const oldProps = instance.props;
+  const oldState = instance.state;
+
+  if (typeof instance.componentWillUpdate === 'function') {
+    instance.componentWillUpdate(props, instance.state, context);
+  }
+
+  instance.context = context;
+  instance.props = props;
+
+  const rendered = instance.render();
+
+  if (typeof instance.componentDidUpdate === 'function') {
+    instance.componentDidUpdate(oldProps, oldState, oldContext);
+  }
+
+  return rendered;
+}
 
 module.exports = ReactShallowRenderer;