From dbadfa2c36f77a8adf1bcae2c85e6f2e0e3d4137 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sema@fb.com>
Date: Tue, 13 Apr 2021 16:57:36 -0400
Subject: [PATCH] [Fizz] Classes Follow Up (#21253)

* Port Classes from Fiber to Fizz

* Test
---
 .../src/__tests__/ReactDOMFizzServer-test.js  | 72 +++++++++++++++++++
 packages/react-server/src/ReactFizzContext.js |  2 +-
 packages/react-server/src/ReactFizzServer.js  |  6 +-
 3 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
index 862c540c3b9f5..9346f20f0442e 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -16,6 +16,7 @@ let React;
 let ReactDOM;
 let ReactDOMFizzServer;
 let Suspense;
+let PropTypes;
 let textCache;
 let document;
 let writable;
@@ -36,6 +37,8 @@ describe('ReactDOMFizzServer', () => {
     }
     Stream = require('stream');
     Suspense = React.Suspense;
+    PropTypes = require('prop-types');
+
     textCache = new Map();
 
     // Test Environment
@@ -655,4 +658,73 @@ describe('ReactDOMFizzServer', () => {
       'http://www.w3.org/2000/svg',
     );
   });
+
+  // @gate experimental
+  it('should can suspend in a class component with legacy context', async () => {
+    class TestProvider extends React.Component {
+      static childContextTypes = {
+        test: PropTypes.string,
+      };
+      state = {ctxToSet: null};
+      static getDerivedStateFromProps(props, state) {
+        return {ctxToSet: props.ctx};
+      }
+      getChildContext() {
+        return {
+          test: this.state.ctxToSet,
+        };
+      }
+      render() {
+        return this.props.children;
+      }
+    }
+
+    class TestConsumer extends React.Component {
+      static contextTypes = {
+        test: PropTypes.string,
+      };
+      render() {
+        const child = (
+          <b>
+            <Text text={this.context.test} />
+          </b>
+        );
+        if (this.props.prefix) {
+          return [readText(this.props.prefix), child];
+        }
+        return child;
+      }
+    }
+
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <TestProvider ctx="A">
+          <div>
+            <Suspense fallback={[<Text text="Loading: " />, <TestConsumer />]}>
+              <TestProvider ctx="B">
+                <TestConsumer prefix="Hello: " />
+              </TestProvider>
+              <TestConsumer />
+            </Suspense>
+          </div>
+        </TestProvider>,
+        writable,
+      );
+      startWriting();
+    });
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        Loading: <b>A</b>
+      </div>,
+    );
+    await act(async () => {
+      resolveText('Hello: ');
+    });
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        Hello: <b>B</b>
+        <b>A</b>
+      </div>,
+    );
+  });
 });
diff --git a/packages/react-server/src/ReactFizzContext.js b/packages/react-server/src/ReactFizzContext.js
index 85751e779d11f..7837d15753d49 100644
--- a/packages/react-server/src/ReactFizzContext.js
+++ b/packages/react-server/src/ReactFizzContext.js
@@ -47,8 +47,8 @@ export function getMaskedContext(type: any, unmaskedContext: Object): Object {
 }
 
 export function processChildContext(
-  type: any,
   instance: any,
+  type: any,
   parentContext: Object,
   childContextTypes: Object,
 ): Object {
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 3664e7ba48830..a3b1babdf5463 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -461,7 +461,7 @@ function renderWithHooks<Props, SecondArg>(
 function finishClassComponent(
   request: Request,
   task: Task,
-  instance: Object,
+  instance: any,
   Component: any,
   props: any,
 ): ReactNodeList {
@@ -518,7 +518,7 @@ function renderClassComponent(
     : undefined;
   const instance = constructClassInstance(Component, props, unmaskedContext);
   mountClassInstance(instance, Component, props, unmaskedContext);
-  finishClassComponent(request, task, Component);
+  finishClassComponent(request, task, instance, Component, props);
 }
 
 const didWarnAboutBadClass = {};
@@ -617,7 +617,7 @@ function renderIndeterminateComponent(
     }
 
     mountClassInstance(value, Component, props, legacyContext);
-    finishClassComponent(request, task, value, Component);
+    finishClassComponent(request, task, value, Component, props);
   } else {
     // Proceed under the assumption that this is a function component
     if (__DEV__) {