From d2f48f91d97182e2abaaa8ba00fc849aa6350e20 Mon Sep 17 00:00:00 2001
From: Kevin Read <me@kevin-read.com>
Date: Sun, 14 Mar 2021 14:55:31 +0100
Subject: [PATCH] Add support for passing context to React class based
 components that request context via setting .contextType, according to
 patches posted in
 https://github.com/enzymejs/enzyme/issues/2189#issuecomment-796416083. Adds
 changes to ReactSixteenAdapter and simple test case.

---
 .../src/ReactSixteenAdapter.js                |  8 +++++-
 .../test/ShallowWrapper-spec.jsx              | 28 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
index 793b0c33e..3a30327ec 100644
--- a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
+++ b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
@@ -690,7 +690,13 @@ class ReactSixteenAdapter extends EnzymeAdapter {
           renderedEl = transformSuspense(renderedEl, renderedEl, { suspenseFallback });
           const { type: Component } = renderedEl;
 
-          const context = getMaskedContext(Component.contextTypes, unmaskedContext);
+          let context;
+          if (Component.contextType) {
+            const Provider = adapter.getProviderFromConsumer(Component.contextType);
+            context = providerValues.has(Provider) ? providerValues.get(Provider) : getProviderDefaultValue(Provider);
+          } else {
+            context = getMaskedContext(Component.contextTypes, unmaskedContext);
+          }
 
           if (isMemo(el.type)) {
             const { type: InnerComp, compare } = el.type;
diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
index a46c62b56..c240c6113 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -666,6 +666,34 @@ describe('shallow', () => {
           expect(consumer.text()).to.equal('foo');
         });
       });
+
+      describe('shallow() on Provider and Consumer through .contextType', () => {
+
+        const { Provider } = React.createContext('howdy!');
+
+        class OuterComponent extends React.Component {
+          render() {
+            return (
+              <Provider value="foo"><InnerComponent /></Provider>
+            );
+          }
+        }
+
+        class InnerComponent extends React.Component {
+          render() {
+            return this.context;
+          }
+        }
+
+        InnerComponent.contextType = Provider;
+
+        it('works on a Provider', () => {
+          const wrapper = shallow(<OuterComponent />);
+          const provides = wrapper.find(Provider).dive();
+          const provider = provides.find(InnerComponent).shallow();
+          expect(provider.text()).to.equal('foo');
+        });
+      });
     });
 
     describeIf(is('> 0.13'), 'stateless function components (SFCs)', () => {