From fc01e4aa7f47ce972c8413d18fa4cc459f46d869 Mon Sep 17 00:00:00 2001
From: Toru Kobayashi <koba0004@gmail.com>
Date: Tue, 26 Dec 2017 23:47:32 +0900
Subject: [PATCH] ShallowRenderer should filter context by contextTypes

---
 .../src/ReactShallowRenderer.js               | 23 +++++++++++++-----
 .../__tests__/ReactShallowRenderer-test.js    | 24 +++++++++++++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/packages/react-test-renderer/src/ReactShallowRenderer.js b/packages/react-test-renderer/src/ReactShallowRenderer.js
index ef50aafa28d1e..8667573d37bb5 100644
--- a/packages/react-test-renderer/src/ReactShallowRenderer.js
+++ b/packages/react-test-renderer/src/ReactShallowRenderer.js
@@ -70,15 +70,15 @@ class ReactShallowRenderer {
 
     this._rendering = true;
     this._element = element;
-    this._context = context;
+    this._context = getMaskedContext(element.type.contextTypes, context);
 
     if (this._instance) {
-      this._updateClassComponent(element.type, element.props, context);
+      this._updateClassComponent(element.type, element.props, this._context);
     } else {
       if (shouldConstruct(element.type)) {
         this._instance = new element.type(
           element.props,
-          context,
+          this._context,
           this._updater,
         );
 
@@ -87,7 +87,7 @@ class ReactShallowRenderer {
 
           checkPropTypes(
             element.type.contextTypes,
-            context,
+            this._context,
             'context',
             getName(element.type, this._instance),
             getStackAddendum,
@@ -96,9 +96,9 @@ class ReactShallowRenderer {
           currentlyValidatingElement = null;
         }
 
-        this._mountClassComponent(element.props, context);
+        this._mountClassComponent(element.props, this._context);
       } else {
-        this._rendered = element.type(element.props, context);
+        this._rendered = element.type(element.props, this._context);
       }
     }
 
@@ -290,4 +290,15 @@ function shouldConstruct(Component) {
   return !!(Component.prototype && Component.prototype.isReactComponent);
 }
 
+function getMaskedContext(contextTypes, unmaskedContext) {
+  if (!contextTypes) {
+    return emptyObject;
+  }
+  const context = {};
+  for (let key in contextTypes) {
+    context[key] = unmaskedContext[key];
+  }
+  return context;
+}
+
 export default ReactShallowRenderer;
diff --git a/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js b/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js
index a551bc35ab33a..d83121e36ae74 100644
--- a/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js
+++ b/packages/react-test-renderer/src/__tests__/ReactShallowRenderer-test.js
@@ -215,6 +215,9 @@ describe('ReactShallowRenderer', () => {
         </div>
       );
     }
+    SomeComponent.contextTypes = {
+      bar: PropTypes.string,
+    };
 
     const shallowRenderer = createRenderer();
     const result = shallowRenderer.render(<SomeComponent foo={'FOO'} />, {
@@ -427,6 +430,9 @@ describe('ReactShallowRenderer', () => {
         super(props, context);
         this.state = initialState;
       }
+      static contextTypes = {
+        context: PropTypes.string,
+      };
       componentDidUpdate(...args) {
         componentDidUpdateParams.push(...args);
       }
@@ -769,6 +775,24 @@ describe('ReactShallowRenderer', () => {
     expect(result).toEqual(<div>foo:baz</div>);
   });
 
+  it('should filter context by contextTypes', () => {
+    class SimpleComponent extends React.Component {
+      static contextTypes = {
+        foo: PropTypes.string,
+      };
+      render() {
+        return <div>{`${this.context.foo}:${this.context.bar}`}</div>;
+      }
+    }
+
+    const shallowRenderer = createRenderer();
+    let result = shallowRenderer.render(<SimpleComponent />, {
+      foo: 'foo',
+      bar: 'bar',
+    });
+    expect(result).toEqual(<div>foo:undefined</div>);
+  });
+
   it('can fail context when shallowly rendering', () => {
     spyOnDev(console, 'error');