[Fizz] Add proper assertion for stream fix (#27543)

gnoff · web-flow · commit 9617d39eca1f · 2023-10-18T12:01:09.000-07:00
In #27541 I added tests to assert that the write after close bug was fixed. However the Node implementation has an abort behavior preventing reproduction of the original issue and the Browser build does not use AsyncLocalStorage which also prevents reproduction. This change moves the Browser test to a Edge runtime where AsyncLocalStorage is polyfilled. In this implementation the test does correctly fail when the patch is removed.
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js
@@ -15,15 +15,13 @@ global.ReadableStream =
 global.TextEncoder = require('util').TextEncoder;
 
 let React;
-let ReactDOM;
 let ReactDOMFizzServer;
 let Suspense;
 
 describe('ReactDOMFizzServerBrowser', () => {
   beforeEach(() => {
     jest.resetModules();
     React = require('react');
-    ReactDOM = require('react-dom');
     ReactDOMFizzServer = require('react-dom/server.browser');
     Suspense = React.Suspense;
   });
@@ -549,37 +547,4 @@ describe('ReactDOMFizzServerBrowser', () => {
     // However, it does error the shell.
     expect(caughtError.message).toEqual('testing postpone');
   });
-
-  // https://github.com/facebook/react/issues/27540
-  // This test is not actually asserting much because in our test environment the Float method cannot find the request after
-  // an await and thus is a noop. If we fix our test environment to support AsyncLocalStorage we can assert that the
-  // stream does not write after closing.
-  it('does not try to write to the stream after it has been closed', async () => {
-    async function preloadLate() {
-      await 1;
-      ReactDOM.preconnect('foo');
-    }
-
-    function Preload() {
-      preloadLate();
-      return null;
-    }
-
-    function App() {
-      return (
-        <html>
-          <body>
-            <main>hello</main>
-            <Preload />
-          </body>
-        </html>
-      );
-    }
-    const stream = await ReactDOMFizzServer.renderToReadableStream(<App />);
-    const result = await readResult(stream);
-
-    expect(result).toMatchInlineSnapshot(
-      `"<!DOCTYPE html><html><head></head><body><main>hello</main></body></html>"`,
-    );
-  });
 });
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerEdge-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerEdge-test.js
@@ -0,0 +1,78 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ */
+
+'use strict';
+
+// Polyfills for test environment
+global.ReadableStream =
+  require('web-streams-polyfill/ponyfill/es6').ReadableStream;
+global.TextEncoder = require('util').TextEncoder;
+global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;
+
+let React;
+let ReactDOM;
+let ReactDOMFizzServer;
+
+describe('ReactDOMFizzServerEdge', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    jest.useRealTimers();
+    React = require('react');
+    ReactDOM = require('react-dom');
+    ReactDOMFizzServer = require('react-dom/server.edge');
+  });
+
+  async function readResult(stream) {
+    const reader = stream.getReader();
+    let result = '';
+    while (true) {
+      const {done, value} = await reader.read();
+      if (done) {
+        return result;
+      }
+      result += Buffer.from(value).toString('utf8');
+    }
+  }
+
+  // https://github.com/facebook/react/issues/27540
+  it('does not try to write to the stream after it has been closed', async () => {
+    async function preloadLate() {
+      await 1;
+      await 1;
+      // need to wait a few microtasks to get the stream to close before this is called
+      ReactDOM.preconnect('foo');
+    }
+
+    function Preload() {
+      preloadLate();
+      return null;
+    }
+
+    function App() {
+      return (
+        <html>
+          <body>
+            <main>hello</main>
+            <Preload />
+          </body>
+        </html>
+      );
+    }
+    const stream = await ReactDOMFizzServer.renderToReadableStream(<App />);
+    const result = await readResult(stream);
+    // need to wait a macrotask to let the scheduled work from the preconnect to execute
+    await new Promise(resolve => {
+      setTimeout(resolve, 1);
+    });
+
+    expect(result).toMatchInlineSnapshot(
+      `"<!DOCTYPE html><html><head></head><body><main>hello</main></body></html>"`,
+    );
+  });
+});
