From ab22518fcc7c5db58f69ee6913869cca01ba66a9 Mon Sep 17 00:00:00 2001
From: Glen Mailer <glenjamin@gmail.com>
Date: Wed, 11 May 2016 09:26:42 +0100
Subject: [PATCH] Allow supplying an error catcher component

---
 src/createClassProxy.js     |  8 ++---
 src/createPrototypeProxy.js | 22 +++++++++++++-
 test/consistency.js         | 58 ++++++++++++++++++++++++++++++++++---
 3 files changed, 79 insertions(+), 9 deletions(-)

diff --git a/src/createClassProxy.js b/src/createClassProxy.js
index c94a2af..b691c16 100644
--- a/src/createClassProxy.js
+++ b/src/createClassProxy.js
@@ -47,7 +47,7 @@ function addProxy(Component, proxy) {
   allProxies.push([Component, proxy]);
 }
 
-function proxyClass(InitialComponent) {
+function proxyClass(InitialComponent, ErrorComponent) {
   // Prevent double wrapping.
   // Given a proxy class, return the existing proxy managing it.
   var existingProxy = findProxy(InitialComponent);
@@ -103,7 +103,7 @@ function proxyClass(InitialComponent) {
   let prototypeProxy;
   if (InitialComponent.prototype && InitialComponent.prototype.isReactComponent) {
     // Point proxy constructor to the proxy prototype
-    prototypeProxy = createPrototypeProxy();
+    prototypeProxy = createPrototypeProxy(ErrorComponent);
     ProxyComponent.prototype = prototypeProxy.get();
   }
 
@@ -257,8 +257,8 @@ function createFallback(Component) {
   };
 }
 
-export default function createClassProxy(Component) {
+export default function createClassProxy(Component, ErrorComponent) {
   return Component.__proto__ && supportsProtoAssignment() ?
-    proxyClass(Component) :
+    proxyClass(Component, ErrorComponent) :
     createFallback(Component);
 }
diff --git a/src/createPrototypeProxy.js b/src/createPrototypeProxy.js
index db2ed7d..e0bd044 100644
--- a/src/createPrototypeProxy.js
+++ b/src/createPrototypeProxy.js
@@ -1,7 +1,8 @@
+import React from "react";
 import assign from 'lodash/assign';
 import difference from 'lodash/difference';
 
-export default function createPrototypeProxy() {
+export default function createPrototypeProxy(ErrorComponent) {
   let proxy = {};
   let current = null;
   let mountedInstances = [];
@@ -43,6 +44,20 @@ export default function createPrototypeProxy() {
     return proxiedMethod;
   }
 
+  /**
+   * Augments the original render with exception handling & rendering.
+   */
+  function proxiedRender() {
+    if (typeof current.render === 'function') {
+      try {
+        return current.render.apply(this, arguments);
+      } catch (ex) {
+        return React.createElement(ErrorComponent, { error: ex });
+      }
+    }
+  }
+  proxiedRender.toString = proxyToString('render');
+
   /**
    * Augments the original componentDidMount with instance tracking.
    */
@@ -158,6 +173,11 @@ export default function createPrototypeProxy() {
       }
     });
 
+    // Error catching for render
+    if (typeof ErrorComponent !== 'undefined') {
+      defineProxyPropertyWithValue('render', proxiedRender);
+    }
+
     // Track mounting and unmounting
     defineProxyPropertyWithValue('componentDidMount', proxiedComponentDidMount);
     defineProxyPropertyWithValue('componentWillUnmount', proxiedComponentWillUnmount);
diff --git a/test/consistency.js b/test/consistency.js
index 6f3a7d1..61bfe39 100644
--- a/test/consistency.js
+++ b/test/consistency.js
@@ -51,7 +51,13 @@ function createModernFixtures() {
   }
   delete Anon.name;
 
-  return { Bar, Baz, Foo, Anon };
+  class Broken extends React.Component {
+    render() {
+      throw new Error('mistakes were made');
+    }
+  }
+
+  return { Bar, Baz, Foo, Anon, Broken };
 }
 
 function createClassicFixtures() {
@@ -101,7 +107,13 @@ function createClassicFixtures() {
   });
   delete Anon.displayName;
 
-  return { Bar, Baz, Foo, Anon };
+  const Broken = React.createClass({
+    render() {
+      throw new Error('mistakes were made');
+    }
+  });
+
+  return { Bar, Baz, Foo, Anon, Broken };
 }
 
 describe('consistency', () => {
@@ -119,9 +131,9 @@ describe('consistency', () => {
   });
 
   function runCommonTests(createFixtures) {
-    let Bar, Baz, Foo, Anon;
+    let Bar, Baz, Foo, Anon, Broken;
     beforeEach(() => {
-      ({ Foo, Bar, Baz, Anon } = createFixtures());
+      ({ Foo, Bar, Baz, Anon, Broken } = createFixtures());
     });
 
     it('does not overwrite the original class', () => {
@@ -284,6 +296,44 @@ describe('consistency', () => {
       const Proxy = proxy.get();
       expect(() => renderer.render(<Proxy />)).toThrow('Oops');
     });
+
+    it('can use error component for catching errors', () => {
+      function ErrorHandler({ error }) {
+        return <p>{error.message}</p>;
+      }
+      let proxy = createProxy(Broken, ErrorHandler);
+      const Proxy = proxy.get();
+      expect(() => renderer.render(<Proxy />)).toNotThrow();
+      const output = renderer.getRenderOutput()
+      expect(output.type).toEqual(ErrorHandler);
+      expect(output.props.error).toBeAn(Error);
+      expect(output.props.error.message).toEqual('mistakes were made');
+    });
+
+    it('still throws errors if no error component supplied', () => {
+      let proxy = createProxy(Broken);
+      const Proxy = proxy.get();
+      expect(() => renderer.render(<Proxy />)).toThrow();
+    });
+
+    it('can recover after handling an error', () => {
+      function ErrorHandler({ error }) {
+        return <p>{error.message}</p>;
+      }
+      let proxy = createProxy(Bar, ErrorHandler);
+      const Proxy = proxy.get();
+
+      expect(() => renderer.render(<Proxy />)).toNotThrow();
+      expect(renderer.getRenderOutput().props.children).toEqual('Bar');
+
+      proxy.update(Broken);
+      expect(() => renderer.render(<Proxy />)).toNotThrow();
+      expect(renderer.getRenderOutput().type).toEqual(ErrorHandler);
+
+      proxy.update(Baz);
+      expect(() => renderer.render(<Proxy />)).toNotThrow();
+      expect(renderer.getRenderOutput().props.children).toEqual('Baz');
+    });
   }
 
   describe('classic', () => {