From 8fcf010483ff90e518cb7885bca5431333247f69 Mon Sep 17 00:00:00 2001
From: Josh Minzner <josh@trialspark.com>
Date: Tue, 22 Jan 2019 14:51:17 -0500
Subject: [PATCH] [new] `shallow`: Support rendering and `dive()`ing
 `createContext()` providers and consumers

---
 .../test/ShallowWrapper-spec.jsx              | 201 ++++++++++++++++++
 .../shared/methods/getWrappingComponent.jsx   |  26 +++
 packages/enzyme/src/ShallowWrapper.js         |  66 ++++--
 3 files changed, 279 insertions(+), 14 deletions(-)

diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
index a999689be..f1659e249 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -270,6 +270,55 @@ describe('shallow', () => {
         expect(wrapper.text()).to.equal('Context says: I can be set!');
       });
 
+      describeIf(is('>= 16.3'), 'with createContext()', () => {
+        let Context1;
+        let Context2;
+
+        function WrappingComponent(props) {
+          const { value1, value2, children } = props;
+          return (
+            <Context1.Provider value={value1}>
+              <Context2.Provider value={value2}>
+                {children}
+              </Context2.Provider>
+            </Context1.Provider>
+          );
+        }
+
+        function Component() {
+          return (
+            <Context1.Consumer>
+              {value1 => (
+                <Context2.Consumer>
+                  {value2 => (
+                    <div>Value 1: {value1}; Value 2: {value2}</div>
+                  )}
+                </Context2.Consumer>
+              )}
+            </Context1.Consumer>
+          );
+        }
+
+        beforeEach(() => {
+          Context1 = createContext('default1');
+          Context2 = createContext('default2');
+        });
+
+        it('renders', () => {
+          const wrapper = shallow(<Component />, {
+            wrappingComponent: WrappingComponent,
+            wrappingComponentProps: {
+              value1: 'one',
+              value2: 'two',
+            },
+          });
+          const consumer1 = wrapper.find(Context1.Consumer).dive();
+          const consumer2 = consumer1.find(Context2.Consumer).dive();
+
+          expect(consumer2.text()).to.equal('Value 1: one; Value 2: two');
+        });
+      });
+
       it('throws an error if the wrappingComponent does not render its children', () => {
         class BadWrapper extends React.Component {
           render() {
@@ -384,6 +433,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 (SFCs)', () => {
       it('can pass in context', () => {
         const SimpleComponent = (props, context) => (
diff --git a/packages/enzyme-test-suite/test/shared/methods/getWrappingComponent.jsx b/packages/enzyme-test-suite/test/shared/methods/getWrappingComponent.jsx
index b90555c1a..7722df0ae 100644
--- a/packages/enzyme-test-suite/test/shared/methods/getWrappingComponent.jsx
+++ b/packages/enzyme-test-suite/test/shared/methods/getWrappingComponent.jsx
@@ -183,6 +183,32 @@ export default function describeGetWrappingComponent({
       }
     });
 
+    itIf(is('>= 16.3'), 'updates a <Provider /> if it is rendered as root', () => {
+      const Context = React.createContext();
+      function WrappingComponent(props) {
+        const { value, children } = props;
+        return (
+          <Context.Provider value={value}>
+            {children}
+          </Context.Provider>
+        );
+      }
+      const wrapper = Wrap(
+        <Context.Consumer>
+          {value => <div>{value}</div>}
+        </Context.Consumer>,
+        {
+          wrappingComponent: WrappingComponent,
+          wrappingComponentProps: { value: 'hello!' },
+        },
+      );
+      const wrappingComponent = wrapper.getWrappingComponent();
+      expect(wrapper.text()).to.equal('hello!');
+
+      wrappingComponent.setProps({ value: 'goodbye!' });
+      expect(wrapper.text()).to.equal('goodbye!');
+    });
+
     itIf(!isShallow, 'handles a partial prop update', () => {
       const wrapper = Wrap(<MyComponent />, {
         wrappingComponent: MyWrappingComponent,
diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
index 4058847fc..6672c1831 100644
--- a/packages/enzyme/src/ShallowWrapper.js
+++ b/packages/enzyme/src/ShallowWrapper.js
@@ -45,6 +45,7 @@ const CHILD_CONTEXT = sym('__childContext__');
 const WRAPPING_COMPONENT = sym('__wrappingComponent__');
 const PRIMARY_WRAPPER = sym('__primaryWrapper__');
 const ROOT_FINDER = sym('__rootFinder__');
+const PROVIDER_VALUES = sym('__providerValues__');
 
 /**
  * Finds all nodes in the current wrapper nodes' render trees that match the provided predicate
@@ -313,14 +314,18 @@ function deepRender(wrapper, target, adapter) {
  * @param {WrappingComponentWrapper} wrapper The `WrappingComponentWrapper` for a
  *  `wrappingComponent`
  * @param {Adapter} adapter An Enzyme adapter
- * @returns {object} The context collected
+ * @returns {object} An object containing an object of legacy context values and a Map of
+ *                   `createContext()` Provider values.
  */
 function getContextFromWrappingComponent(wrapper, adapter) {
   const rootFinder = deepRender(wrapper, wrapper[ROOT_FINDER], adapter);
   if (!rootFinder) {
     throw new Error('`wrappingComponent` must render its children!');
   }
-  return rootFinder[OPTIONS].context;
+  return {
+    legacyContext: rootFinder[OPTIONS].context,
+    providerValues: rootFinder[PROVIDER_VALUES],
+  };
 }
 
 /**
@@ -338,6 +343,7 @@ function getContextFromWrappingComponent(wrapper, adapter) {
 function makeShallowOptions(nodes, root, passedOptions, wrapper) {
   const options = makeOptions(passedOptions);
   const adapter = getAdapter(passedOptions);
+  privateSet(options, PROVIDER_VALUES, passedOptions[PROVIDER_VALUES]);
   if (root || !isCustomComponent(options.wrappingComponent, adapter)) {
     return options;
   }
@@ -347,16 +353,18 @@ function makeShallowOptions(nodes, root, passedOptions, wrapper) {
   const { node: wrappedNode, RootFinder } = adapter.wrapWithWrappingComponent(nodes, options);
   // eslint-disable-next-line no-use-before-define
   const wrappingComponent = new WrappingComponentWrapper(wrappedNode, wrapper, RootFinder);
-  const wrappingComponentContext = getContextFromWrappingComponent(
-    wrappingComponent, adapter,
-  );
+  const {
+    legacyContext: wrappingComponentLegacyContext,
+    providerValues: wrappingComponentProviderValues,
+  } = getContextFromWrappingComponent(wrappingComponent, adapter);
   privateSet(wrapper, WRAPPING_COMPONENT, wrappingComponent);
   return {
     ...options,
     context: {
       ...options.context,
-      ...wrappingComponentContext,
+      ...wrappingComponentLegacyContext,
     },
+    [PROVIDER_VALUES]: wrappingComponentProviderValues,
   };
 }
 
@@ -385,10 +393,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(options[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) {
@@ -415,6 +425,7 @@ class ShallowWrapper {
       privateSetNodes(this, nodes);
       privateSet(this, OPTIONS, root[OPTIONS]);
       privateSet(this, ROOT_NODES, root[NODES]);
+      privateSet(this, PROVIDER_VALUES, null);
     }
   }
 
@@ -612,7 +623,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();
@@ -1659,14 +1672,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);
     });
   }
 
@@ -1686,14 +1701,37 @@ class ShallowWrapper {
  * `wrappingComponent` re-renders.
  */
 function updatePrimaryRootContext(wrappingComponent) {
-  const context = getContextFromWrappingComponent(
+  const adapter = getAdapter(wrappingComponent[OPTIONS]);
+  const primaryWrapper = wrappingComponent[PRIMARY_WRAPPER];
+  const primaryRenderer = primaryWrapper[RENDERER];
+  const primaryNode = primaryRenderer.getNode();
+  const primaryNodeIsContextConsumer = typeof adapter.isContextConsumer === 'function'
+    && adapter.isContextConsumer(primaryNode.type);
+  const { legacyContext, providerValues } = getContextFromWrappingComponent(
     wrappingComponent,
-    getAdapter(wrappingComponent[OPTIONS]),
+    adapter,
   );
-  wrappingComponent[PRIMARY_WRAPPER].setContext({
+  const prevProviderValues = primaryWrapper[PROVIDER_VALUES];
+
+  primaryWrapper.setContext({
     ...wrappingComponent[PRIMARY_WRAPPER][OPTIONS].context,
-    ...context,
+    ...legacyContext,
   });
+  primaryWrapper[PROVIDER_VALUES] = new Map([...prevProviderValues, ...providerValues]);
+
+  if (primaryNodeIsContextConsumer) {
+    const Consumer = primaryNode.type;
+    // Adapters with an `isContextConsumer` method will definitely have a `getProviderFromConsumer`
+    // method.
+    const Provider = adapter.getProviderFromConsumer(Consumer);
+    const newValue = providerValues.get(Provider);
+    const oldValue = prevProviderValues.get(Provider);
+
+    // Use referential comparison like React
+    if (newValue !== oldValue) {
+      primaryWrapper.rerender();
+    }
+  }
 }
 
 /**