[new] shallow: Support rendering and dive()ing createContext() …

…providers and consumers

packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx

@@ -239,6 +239,158 @@ describe('shallow', () => {
       expect(wrapper.find('.child2')).to.have.lengthOf(1);
     });
 
+    describeIf(is('>= 16.3'), 'createContext()', () => {
+      describe('rendering as root:', () => {
+        let Context;
+
+        beforeEach(() => {
+          Context = createContext('cool');
+        });
+
+        describe('<Provider />', () => {
+          it('can be rendered as the root', () => {
+            const wrapper = shallow(
+              <Context.Provider value="hello">
+                <Context.Consumer>
+                  {value => <div>{value}</div>}
+                </Context.Consumer>
+              </Context.Provider>,
+            );
+            expect(wrapper.debug()).to.eql(`
+<ContextConsumer>
+  [function]
+</ContextConsumer>
+            `.trim());
+          });
+
+          it('supports changing the value', () => {
+            const wrapper = shallow(
+              <Context.Provider value="hello">
+                <Context.Consumer>
+                  {value => <div>{value}</div>}
+                </Context.Consumer>
+              </Context.Provider>,
+            );
+            wrapper.setProps({ value: 'world' });
+            expect(wrapper.find(Context.Consumer).dive().text()).to.eql('world');
+          });
+        });
+
+        describe('<Consumer />', () => {
+          function DivRenderer({ children }) {
+            return <div>{children}</div>;
+          }
+          it('can be rendered as the root', () => {
+            const wrapper = shallow(
+              <Context.Consumer>
+                {value => <DivRenderer>{value}</DivRenderer>}
+              </Context.Consumer>,
+            );
+            expect(wrapper.debug()).to.eql(`
+<DivRenderer>
+  cool
+</DivRenderer>
+            `.trim());
+          });
+
+          it('supports changing the children', () => {
+            const wrapper = shallow(
+              <Context.Consumer>
+                {value => <DivRenderer>{value}</DivRenderer>}
+              </Context.Consumer>,
+            );
+            wrapper.setProps({ children: value => <DivRenderer>Changed: {value}</DivRenderer> });
+            expect(wrapper.find(DivRenderer).dive().text()).to.eql('Changed: cool');
+          });
+        });
+      });
+
+      describe('dive() on Provider and Consumer', () => {
+        let Provider;
+        let Consumer;
+
+        beforeEach(() => {
+          ({ Provider, Consumer } = React.createContext('howdy!'));
+        });
+
+        class Consumes extends React.Component {
+          render() {
+            return (
+              <span>
+                <Consumer>{value => <span>{value}</span>}</Consumer>
+              </span>
+            );
+          }
+        }
+
+        class Provides extends React.Component {
+          render() {
+            const { children } = this.props;
+
+            return (
+              <Provider value="foo"><div><div />{children}</div></Provider>
+            );
+          }
+        }
+
+        class MyComponent extends React.Component {
+          render() {
+            return (
+              <Provides><Consumes /></Provides>
+            );
+          }
+        }
+
+        it('works on a Provider', () => {
+          expect(shallow(<MyComponent />)
+            .find(Provides)
+            .dive()
+            .find(Provider)
+            .dive()
+            .text()).to.equal('<Consumes />');
+        });
+
+        it('always gives the default provider value if dive()ing directly to a <Consumer />', () => {
+          // Diving directly on a consumer will give you the default value
+          expect(shallow(<MyComponent />)
+            .find(Consumes)
+            .dive()
+            .find(Consumer)
+            .dive()
+            .text()).to.equal('howdy!');
+        });
+
+        it('gives the actual <Provider /> value if one dive()s it', () => {
+          expect(shallow(<MyComponent />)
+            .find(Provides)
+            .dive()
+            .find(Provider)
+            .dive()
+            .find(Consumes)
+            .dive()
+            .find(Consumer)
+            .dive()
+            .text()).to.equal('foo');
+        });
+
+        it('does not leak values across roots', () => {
+          const wrapper = shallow(<MyComponent />);
+
+          wrapper
+            .find(Provides)
+            .dive()
+            .find(Provider)
+            .dive();
+          expect(wrapper
+            .find(Consumes)
+            .dive()
+            .find(Consumer)
+            .dive()
+            .text()).to.equal('howdy!');
+        });
+      });
+    });
+
     describeIf(is('> 0.13'), 'stateless function components', () => {
       it('can pass in context', () => {
         const SimpleComponent = (props, context) => (

packages/enzyme/src/ShallowWrapper.js

@@ -41,6 +41,7 @@ const OPTIONS = sym('__options__');
 const SET_STATE = sym('__setState__');
 const ROOT_NODES = sym('__rootNodes__');
 const CHILD_CONTEXT = sym('__childContext__');
+const PROVIDER_VALUES = sym('__providerValues__');
 
 /**
  * Finds all nodes in the current wrapper nodes' render trees that match the provided predicate
@@ -268,10 +269,12 @@ class ShallowWrapper {
       privateSet(this, UNRENDERED, nodes);
       const renderer = adapter.createRenderer({ mode: 'shallow', ...options });
       privateSet(this, RENDERER, renderer);
-      this[RENDERER].render(nodes, options.context);
+      const providerValues = new Map(passedOptions[PROVIDER_VALUES] || []);
+      this[RENDERER].render(nodes, options.context, { providerValues });
       const renderedNode = this[RENDERER].getNode();
       privateSetNodes(this, getRootNode(renderedNode));
       privateSet(this, OPTIONS, options);
+      privateSet(this, PROVIDER_VALUES, providerValues);
 
       const { instance } = renderedNode;
       if (instance && !options.disableLifecycleMethods) {
@@ -298,6 +301,7 @@ class ShallowWrapper {
       privateSetNodes(this, nodes);
       privateSet(this, OPTIONS, root[OPTIONS]);
       privateSet(this, ROOT_NODES, root[NODES]);
+      privateSet(this, PROVIDER_VALUES, null);
     }
   }
 
@@ -470,7 +474,9 @@ class ShallowWrapper {
             );
           }
           if (props) this[UNRENDERED] = cloneElement(adapter, this[UNRENDERED], props);
-          this[RENDERER].render(this[UNRENDERED], nextContext);
+          this[RENDERER].render(this[UNRENDERED], nextContext, {
+            providerValues: this[PROVIDER_VALUES],
+          });
           if (shouldComponentUpdateSpy) {
             shouldRender = shouldComponentUpdateSpy.getLastReturnValue();
             shouldComponentUpdateSpy.restore();
@@ -1504,14 +1510,16 @@ class ShallowWrapper {
       if (!isCustomComponentElement(el, adapter)) {
         throw new TypeError(`ShallowWrapper::${name}() can only be called on components`);
       }
-      return this.wrap(el, null, {
+      const childOptions = {
         ...this[OPTIONS],
         ...options,
         context: options.context || {
           ...this[OPTIONS].context,
           ...this[ROOT][CHILD_CONTEXT],
         },
-      });
+      };
+      privateSet(childOptions, PROVIDER_VALUES, this[ROOT][PROVIDER_VALUES]);
+      return this.wrap(el, null, childOptions);
     });
   }