From 16cf05f030da953760c6928b60a7acb3ae5fd821 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] [enzyme-adapter-react-16] [new] add `contextType` support

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.

Co-authored-by: Kevin Read <me@kevin-read.com>
Co-authored-by: Pablo Palacios <pablo.palacios@holidaycheck.com>
---
 .../src/ReactSixteenAdapter.js                |  8 ++-
 .../test/ShallowWrapper-spec.jsx              | 72 +++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
index b07a1f31b..937e1f794 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 efd22c9a4..6f7458da5 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -666,6 +666,78 @@ describe('shallow', () => {
           expect(consumer.text()).to.equal('foo');
         });
       });
+
+      describeIf(is('>= 16.6'), 'shallow() on Provider and Consumer through .contextType', () => {
+        const { Provider } = createContext('howdy!');
+
+        class OuterComponent extends React.Component {
+          render() {
+            const { value } = this.props;
+            return (
+              <Provider value={value}><InnerComponent /></Provider>
+            );
+          }
+        }
+
+        class WrappingComponent extends React.Component {
+          render() {
+            const { children, value } = this.props;
+            return (
+              <Provider value={value}>{ children }</Provider>
+            );
+          }
+        }
+
+        class InnerComponent extends React.Component {
+          render() {
+            return this.context;
+          }
+        }
+
+        InnerComponent.contextType = Provider;
+
+        describe('rendering the Provider directly', () => {
+          it('renders initial context value', () => {
+            const wrapper = shallow(<OuterComponent value="foo" />);
+            const provides = wrapper.find(Provider).shallow();
+            const provider = provides.find(InnerComponent).dive();
+
+            expect(provider.text()).to.equal('foo');
+          });
+
+          it('renders updated context value', () => {
+            const wrapper = shallow(<OuterComponent value="foo" />);
+            wrapper.setProps({ value: 'bar' });
+            const provides = wrapper.find(Provider).shallow();
+            const provider = provides.find(InnerComponent).dive();
+
+            expect(provider.text()).to.equal('bar');
+          });
+        });
+
+        describe('rendering the Provider through wrappingComponent', () => {
+          it('renders initial context value', () => {
+            const wrapper = shallow(<InnerComponent />, {
+              wrappingComponent: WrappingComponent,
+              wrappingComponentProps: { value: 'foo' },
+            });
+
+            expect(wrapper.text()).to.equal('foo');
+          });
+
+          it('renders updated context value', () => {
+            const wrapper = shallow(<InnerComponent />, {
+              wrappingComponent: WrappingComponent,
+              wrappingComponentProps: { value: 'foo' },
+            });
+            const wrappingComponent = wrapper.getWrappingComponent();
+            wrappingComponent.setProps({ value: 'bar' });
+            wrappingComponent.rerender();
+
+            expect(wrapper.text()).to.equal('bar');
+          });
+        });
+      });
     });
 
     describeIf(is('> 0.13'), 'stateless function components (SFCs)', () => {