From b1eb170bce08616fe5de97474605cc709e97043f Mon Sep 17 00:00:00 2001
From: lamhieu-vk <lamhieu.vk@gmail.com>
Date: Fri, 28 Aug 2020 17:09:20 +0700
Subject: [PATCH] [enzyme-adapter-react-16] [fix] `shallow`: missing instance
 when using Suspense

Fixes #2434.
---
 .../src/ReactSixteenAdapter.js                | 42 +++++++++++++++----
 .../test/ShallowWrapper-spec.jsx              | 18 ++++++++
 2 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
index e4d445e5d..99632b839 100644
--- a/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
+++ b/packages/enzyme-adapter-react-16/src/ReactSixteenAdapter.js
@@ -120,20 +120,44 @@ function unmemoType(type) {
   return isMemo(type) ? type.type : type;
 }
 
-function checkIsSuspenseAndCloneElement(el, { suspenseFallback }) {
-  if (!isSuspense(el)) {
-    return el;
+function transformSuspense(renderedEl, prerenderEl, { suspenseFallback }) {
+  if (!isSuspense(renderedEl)) {
+    return renderedEl;
   }
 
-  let { children } = el.props;
+  let { children } = renderedEl.props;
 
   if (suspenseFallback) {
-    const { fallback } = el.props;
+    const { fallback } = renderedEl.props;
     children = replaceLazyWithFallback(children, fallback);
   }
 
-  const FakeSuspenseWrapper = (props) => React.createElement(el.type, { ...el.props, ...props }, children);
-  return React.createElement(FakeSuspenseWrapper, null, children);
+  if (isStateful(prerenderEl.type)) {
+    class FakeSuspense extends prerenderEl.type {
+      render() {
+        return React.createElement(
+          prerenderEl.type,
+          { ...prerenderEl.props, ...this.props },
+          children,
+        );
+      }
+    }
+
+    return React.createElement(FakeSuspense, null, children);
+  }
+
+  return React.createElement(
+    // eslint-disable-next-line prefer-arrow-callback
+    function FakeSuspense(props) {
+      return React.createElement(
+        renderedEl.type,
+        { ...renderedEl.props, ...props },
+        children,
+      );
+    },
+    null,
+    children,
+  );
 }
 
 function elementToTree(el) {
@@ -609,7 +633,7 @@ class ReactSixteenAdapter extends EnzymeAdapter {
 
       const typeIsExisted = !!(renderedEl && renderedEl.type);
       if (is166 && typeIsExisted) {
-        const clonedEl = checkIsSuspenseAndCloneElement(renderedEl, { suspenseFallback });
+        const clonedEl = transformSuspense(renderedEl, elConfig, { suspenseFallback });
 
         const elementIsChanged = clonedEl.type !== renderedEl.type;
         if (elementIsChanged) {
@@ -652,7 +676,7 @@ class ReactSixteenAdapter extends EnzymeAdapter {
             throw TypeError('`React.lazy` is not supported by shallow rendering.');
           }
 
-          renderedEl = checkIsSuspenseAndCloneElement(renderedEl, { suspenseFallback });
+          renderedEl = transformSuspense(renderedEl, renderedEl, { suspenseFallback });
           const { type: Component } = renderedEl;
 
           const context = getMaskedContext(Component.contextTypes, unmaskedContext);
diff --git a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
index 98ab82001..117d64efe 100644
--- a/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
+++ b/packages/enzyme-test-suite/test/ShallowWrapper-spec.jsx
@@ -1930,6 +1930,24 @@ describe('shallow', () => {
       const LazyComponent = lazy(() => fakeDynamicImport(DynamicComponent));
       expect(() => shallow(<LazyComponent />)).to.throw();
     });
+
+    it('returns the correct instance if using Suspense in stateful components', () => {
+      const LazyComponent = lazy(() => fakeDynamicImport(DynamicComponent));
+
+      class Bar extends React.Component {
+        render() {
+          return (
+            <Suspense fallback={<Fallback />}>
+              <LazyComponent />
+            </Suspense>
+          );
+        }
+      }
+
+      const wrapper = shallow(<Bar />);
+
+      expect(wrapper.instance()).to.instanceOf(Bar);
+    });
   });
 
   describe('lifecycle methods', () => {