Add ability to dive() Consumers and Providers

enzymejs · Jan 8, 2019 · e35061c · e35061c
1 parent d66104f
commit e35061c
Show file tree

Hide file tree

Showing 5 changed files with 179 additions and 6 deletions.
diff --git a/docs/api/ShallowWrapper/dive.md b/docs/api/ShallowWrapper/dive.md
@@ -46,3 +46,30 @@ expect(wrapper.find('.in-bar')).to.have.lengthOf(0);
 expect(wrapper.find(Bar)).to.have.lengthOf(1);
 expect(wrapper.find(Bar).dive().find('.in-bar')).to.have.lengthOf(1);
 ```
+
+
+
+#### Common Gotchas
+
+- If you _only_ dive a `<Consumer />` (even if there is a `<Provider />` in your component tree) you always get the default context value. If you want the `<Consumer />` to receive your `<Provider />`'s `value` you must first `.dive()` the `<Provider>`, the `.dive()` the `<Consumer />`.
+  ```jsx
+  const { Provider, Consumer } = React.createContext('foo');
+  class MyComponent extends React.Component {
+    render() {
+      return (
+        <div>
+          <Provider value="bar">
+            <Consumer>
+              {value => (
+                <div>{value}</div>
+              )}
+            </Consumer>
+          </Provider>
+        </div>
+      )
+    }
+  }
+  const wrapper = shallow(<MyComponent />);
+  wrapper.find(Consumer).dive().text(); // "foo"
+  wrapper.find(Provider).dive().find(Consumer).dive().text(); // "bar"
+  ```
diff --git a/packages/enzyme-adapter-react-16.3/src/ReactSixteenThreeAdapter.js b/packages/enzyme-adapter-react-16.3/src/ReactSixteenThreeAdapter.js
@@ -13,6 +13,8 @@ import {
   isPortal,
   isForwardRef,
   isValidElementType,
+  isContextProvider,
+  isContextConsumer,
   AsyncMode,
   Fragment,
   ContextConsumer,
@@ -321,11 +323,21 @@ class ReactSixteenThreeAdapter extends EnzymeAdapter {
     let isDOM = false;
     let cachedNode = null;
     return {
-      render(el, context) {
+      render(el, context, {
+        providerValues = new Map(),
+      } = {}) {
         cachedNode = el;
         /* eslint consistent-return: 0 */
         if (typeof el.type === 'string') {
           isDOM = true;
+        } else if (isContextProvider(el)) {
+          providerValues.set(el.type, el.props.value);
+          return this.render(el.props.children, context, { providerValues });
+        } else if (isContextConsumer(el)) {
+          const value = providerValues.has(el.type.Provider)
+            ? providerValues.get(el.type.Provider)
+            : el.type.Provider._context._defaultValue;
+          return this.render(el.props.children(value), context, { providerValues });
         } else {
           isDOM = false;
           const { type: Component } = el;
@@ -493,7 +505,10 @@ class ReactSixteenThreeAdapter extends EnzymeAdapter {
     if (!inst || !this.isValidElement(inst)) {
       return false;
     }
-    return typeof inst.type === 'function' || isForwardRef(inst);
+    return typeof inst.type === 'function'
+      || isForwardRef(inst)
+      || isContextProvider(inst)
+      || isContextConsumer(inst);
   }
 
   createElement(...args) {

diff --git a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
@@ -13,6 +13,8 @@ import {
   isPortal,
   isForwardRef,
   isValidElementType,
+  isContextProvider,
+  isContextConsumer,
   AsyncMode,
   ConcurrentMode,
   Fragment,
@@ -249,6 +251,31 @@ function getEmptyStateValue() {
   return testRenderer._instance.state;
 }
 
+function getProviderFromConsumer(Consumer) {
+  // React stores references to the Provider on a Consumer differently across versions.
+  let Provider;
+  if (Consumer.Provider) {
+    ({ Provider } = Consumer);
+  } else if (Consumer._context) {
+    ({ Provider } = Consumer._context);
+  }
+  if (Provider) {
+    return Provider;
+  }
+  throw new Error('Enzyme Internal Error: can\'t figure out how to get Provider from Consumer');
+}
+
+function getProviderDefaultValue(Provider) {
+  // React stores references to the Provider's defaultValue differently across versions.
+  if ('_defaultValue' in Provider._context) {
+    return Provider._context._defaultValue;
+  }
+  if ('_currentValue' in Provider._context) {
+    return Provider._context._currentValue;
+  }
+  throw new Error('Enzyme Internal Error: can\'t figure out how to get Provider\'s default value');
+}
+
 class ReactSixteenAdapter extends EnzymeAdapter {
   constructor() {
     super();
@@ -343,11 +370,21 @@ class ReactSixteenAdapter extends EnzymeAdapter {
     let isDOM = false;
     let cachedNode = null;
     return {
-      render(el, context) {
+      render(el, context, {
+        providerValues = new Map(),
+      } = {}) {
         cachedNode = el;
         /* eslint consistent-return: 0 */
         if (typeof el.type === 'string') {
           isDOM = true;
+        } else if (isContextProvider(el)) {
+          providerValues.set(el.type, el.props.value);
+          return this.render(el.props.children, context, { providerValues });
+        } else if (isContextConsumer(el)) {
+          const Provider = getProviderFromConsumer(el.type);
+          const value = providerValues.has(Provider)
+            ? providerValues.get(Provider) : getProviderDefaultValue(Provider);
+          return this.render(el.props.children(value), context, { providerValues });
         } else {
           isDOM = false;
           const { type: Component } = el;
@@ -540,7 +577,10 @@ class ReactSixteenAdapter extends EnzymeAdapter {
     if (!inst || !this.isValidElement(inst)) {
       return false;
     }
-    return typeof inst.type === 'function' || isForwardRef(inst);
+    return typeof inst.type === 'function'
+      || isForwardRef(inst)
+      || isContextProvider(inst)
+      || isContextConsumer(inst);
   }
 
   createElement(...args) {

diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -239,6 +239,91 @@ describe('shallow', () => {
       expect(wrapper.find('.child2')).to.have.lengthOf(1);
     });
 
+    describeIf(is('>= 16.3'), '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) => (

diff --git a/packages/enzyme/src/ShallowWrapper.js b/packages/enzyme/src/ShallowWrapper.js
@@ -40,6 +40,7 @@ const ROOT = sym('__root__');
 const OPTIONS = sym('__options__');
 const SET_STATE = sym('__setState__');
 const ROOT_NODES = sym('__rootNodes__');
+const PROVIDER_VALUES = sym('__providerValues__');
 
 /**
  * Finds all nodes in the current wrapper nodes' render trees that match the provided predicate
@@ -201,10 +202,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) {
@@ -230,6 +233,7 @@ class ShallowWrapper {
       privateSetNodes(this, nodes);
       privateSet(this, OPTIONS, root[OPTIONS]);
       privateSet(this, ROOT_NODES, root[NODES]);
+      privateSet(this, PROVIDER_VALUES, null);
     }
   }
 
@@ -1411,7 +1415,9 @@ class ShallowWrapper {
       if (!isCustomComponentElement(el, adapter)) {
         throw new TypeError(`ShallowWrapper::${name}() can only be called on components`);
       }
-      return this.wrap(el, null, { ...this[OPTIONS], ...options });
+      const childOptions = { ...this[OPTIONS], ...options };
+      privateSet(childOptions, PROVIDER_VALUES, this[ROOT][PROVIDER_VALUES]);
+      return this.wrap(el, null, childOptions);
     });
   }