Merge pull request #18737 from storybookjs/fix-react-18-docs-tests

React: Fix callback behavior in `react@18`
storybookjs · Jul 31, 2022 · 1bb8d80 · 1bb8d80
2 parents 0180f5b + b2b7a15
commit 1bb8d80
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 25 deletions.
diff --git a/code/cypress/generated/addon-docs.spec.ts b/code/cypress/generated/addon-docs.spec.ts
@@ -8,28 +8,22 @@ describe('addon-docs', () => {
 
   skipOn('vue3', () => {
     skipOn('html', () => {
-      skipOn('react', () => {
-        skipOn('cra', () => {
-          skipOn('react_legacy_root_api', () => {
-            it('should provide source snippet', () => {
-              cy.getDocsElement()
-                .find('.docblock-code-toggle')
-                .each(($div) => {
-                  cy.wrap($div)
-                    .should('contain.text', 'Show code')
-                    // use force click so cypress does not automatically scroll, making the source block visible on this step
-                    .click({ force: true });
-                });
+      it('should provide source snippet', () => {
+        cy.getDocsElement()
+          .find('.docblock-code-toggle')
+          .each(($div) => {
+            cy.wrap($div)
+              .should('contain.text', 'Show code')
+              // use force click so cypress does not automatically scroll, making the source block visible on this step
+              .click({ force: true });
+          });
 
-              cy.getDocsElement()
-                .find('pre.prismjs')
-                .each(($div) => {
-                  const text = $div.text();
-                  expect(text).not.match(/^\(args\) => /);
-                });
-            });
+        cy.getDocsElement()
+          .find('pre.prismjs')
+          .each(($div) => {
+            const text = $div.text();
+            expect(text).not.match(/^\(args\) => /);
           });
-        });
       });
     });
   });

diff --git a/code/renderers/react/src/render.tsx b/code/renderers/react/src/render.tsx
@@ -1,7 +1,15 @@
 // @ts-ignore
 import global from 'global';
 
-import React, { Component as ReactComponent, FC, ReactElement, StrictMode, Fragment } from 'react';
+import React, {
+  Component as ReactComponent,
+  FC,
+  ReactElement,
+  StrictMode,
+  Fragment,
+  useLayoutEffect,
+  useRef,
+} from 'react';
 import ReactDOM, { version as reactDomVersion } from 'react-dom';
 import type { Root as ReactRoot } from 'react-dom/client';
 
@@ -26,16 +34,32 @@ export const render: ArgsStoryFn<ReactFramework> = (args, context) => {
   return <Component {...args} />;
 };
 
+const WithCallback: FC<{ callback: () => void; children: ReactElement }> = ({
+  callback,
+  children,
+}) => {
+  // See https://github.com/reactwg/react-18/discussions/5#discussioncomment-2276079
+  const once = useRef(false);
+  useLayoutEffect(() => {
+    if (once.current) return;
+    once.current = true;
+    callback();
+  }, [callback]);
+
+  return children;
+};
+
 const renderElement = async (node: ReactElement, el: Element) => {
   // Create Root Element conditionally for new React 18 Root Api
   const root = await getReactRoot(el);
 
   return new Promise((resolve) => {
     if (root) {
-      root.render(node);
-      setTimeout(() => {
-        resolve(null);
-      }, 0);
+      root.render(
+        <WithCallback key={Math.random()} callback={() => resolve(null)}>
+          {node}
+        </WithCallback>
+      );
     } else {
       ReactDOM.render(node, el, () => resolve(null));
     }