From c69b712d0a42c8a863bd1cf7dee3042ca5d7002c Mon Sep 17 00:00:00 2001
From: Josh Minzner <josh@trialspark.com>
Date: Tue, 22 Jan 2019 14:50:13 -0500
Subject: [PATCH] [enzyme-adapter-react-{16,16.3}] [new] support shallow
 rendering `createContext()` providers and consumers

---
 .../src/ReactSixteenThreeAdapter.js           | 24 ++++++++-
 .../src/ReactSixteenAdapter.js                | 49 ++++++++++++++++++-
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/packages/enzyme-adapter-react-16.3/src/ReactSixteenThreeAdapter.js b/packages/enzyme-adapter-react-16.3/src/ReactSixteenThreeAdapter.js
index 74a209753..6266d46dd 100644
--- a/packages/enzyme-adapter-react-16.3/src/ReactSixteenThreeAdapter.js
+++ b/packages/enzyme-adapter-react-16.3/src/ReactSixteenThreeAdapter.js
@@ -15,6 +15,8 @@ import {
   Element,
   ForwardRef,
   Fragment,
+  isContextConsumer,
+  isContextProvider,
   isElement,
   isForwardRef,
   isPortal,
@@ -357,11 +359,29 @@ 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);
+          const MockProvider = Object.assign(
+            props => props.children,
+            el.type,
+          );
+          return withSetStateAllowed(() => renderer.render({ ...el, type: MockProvider }));
+        } else if (isContextConsumer(el)) {
+          const value = providerValues.has(el.type.Provider)
+            ? providerValues.get(el.type.Provider)
+            : el.type.Provider._context._defaultValue;
+          const MockConsumer = Object.assign(
+            props => props.children(value),
+            el.type,
+          );
+          return withSetStateAllowed(() => renderer.render({ ...el, type: MockConsumer }));
         } else {
           isDOM = false;
           const { type: Component } = el;
@@ -543,6 +563,8 @@ class ReactSixteenThreeAdapter extends EnzymeAdapter {
     return !!type && (
       typeof type === 'function'
       || isForwardRef(fakeElement)
+      || isContextProvider(fakeElement)
+      || isContextConsumer(fakeElement)
     );
   }
 
diff --git a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
index cbe6750fc..c64c67572 100644
--- a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
+++ b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
@@ -18,6 +18,8 @@ import {
   Element,
   ForwardRef,
   Fragment,
+  isContextConsumer,
+  isContextProvider,
   isElement,
   isForwardRef,
   isMemo,
@@ -298,6 +300,31 @@ function wrapAct(fn) {
   return returnVal;
 }
 
+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();
@@ -449,11 +476,29 @@ class ReactSixteenAdapter extends EnzymeAdapter {
     };
 
     return {
-      render(el, unmaskedContext) {
+      render(el, unmaskedContext, {
+        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);
+          const MockProvider = Object.assign(
+            props => props.children,
+            el.type,
+          );
+          return withSetStateAllowed(() => renderer.render({ ...el, type: MockProvider }));
+        } else if (isContextConsumer(el)) {
+          const Provider = getProviderFromConsumer(el.type);
+          const value = providerValues.has(Provider)
+            ? providerValues.get(Provider) : getProviderDefaultValue(Provider);
+          const MockConsumer = Object.assign(
+            props => props.children(value),
+            el.type,
+          );
+          return withSetStateAllowed(() => renderer.render({ ...el, type: MockConsumer }));
         } else {
           isDOM = false;
           const { type: Component } = el;
@@ -672,6 +717,8 @@ class ReactSixteenAdapter extends EnzymeAdapter {
     return !!type && (
       typeof type === 'function'
       || isForwardRef(fakeElement)
+      || isContextProvider(fakeElement)
+      || isContextConsumer(fakeElement)
     );
   }