React: Use Act wrapper in Storybook for component rendering

storybookjs · Dec 12, 2024 · b4a1714 · b4a1714
1 parent 47cbf5a
commit b4a1714
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 17 deletions.
diff --git a/code/renderers/react/src/act-compat.ts b/code/renderers/react/src/act-compat.ts
@@ -26,7 +26,7 @@ export function getReactActEnvironment() {
 }
 
 function withGlobalActEnvironment(actImplementation: (callback: () => void) => Promise<any>) {
-  return (callback: () => any) => {
+  return async (callback: () => any) => {
     const previousActEnvironment = getReactActEnvironment();
     setReactActEnvironment(true);
     try {
@@ -40,8 +40,8 @@ function withGlobalActEnvironment(actImplementation: (callback: () => void) => P
         return result;
       });
       if (callbackNeedsToBeAwaited) {
-        const thenable: Promise<any> = actResult;
-        return {
+        const thenable = actResult;
+        return await {
           then: (resolve: (param: any) => void, reject: (param: any) => void) => {
             thenable.then(
               (returnValue) => {
@@ -57,7 +57,7 @@ function withGlobalActEnvironment(actImplementation: (callback: () => void) => P
         };
       } else {
         setReactActEnvironment(previousActEnvironment);
-        return actResult;
+        return await actResult;
       }
     } catch (error) {
       // Can't be a `finally {}` block since we don't know if we have to immediately restore IS_REACT_ACT_ENVIRONMENT

diff --git a/code/renderers/react/src/portable-stories.tsx b/code/renderers/react/src/portable-stories.tsx
@@ -106,16 +106,11 @@ export const INTERNAL_DEFAULT_PROJECT_ANNOTATIONS: ProjectAnnotations<ReactRende
   },
   renderToCanvas: async (renderContext, canvasElement) => {
     if (renderContext.storyContext.testingLibraryRender == null) {
-      let unmount: () => void;
-
-      await act(async () => {
-        unmount = await reactProjectAnnotations.renderToCanvas(renderContext, canvasElement);
-      });
+      renderContext.storyContext.globals.isPortableStory = true;
+      const unmount = await reactProjectAnnotations.renderToCanvas(renderContext, canvasElement);
 
       return async () => {
-        await act(() => {
-          unmount();
-        });
+        await unmount();
       };
     }
     const {

diff --git a/code/renderers/react/src/renderToCanvas.tsx b/code/renderers/react/src/renderToCanvas.tsx
@@ -5,7 +5,7 @@ import type { RenderContext } from 'storybook/internal/types';
 
 import { global } from '@storybook/global';
 
-import { getReactActEnvironment } from './act-compat';
+import { act } from './act-compat';
 import type { ReactRenderer, StoryContext } from './types';
 
 const { FRAMEWORK_OPTIONS } = global;
@@ -58,9 +58,10 @@ export async function renderToCanvas(
   const { renderElement, unmountElement } = await import('@storybook/react-dom-shim');
   const Story = unboundStoryFn as FC<StoryContext<ReactRenderer>>;
 
-  const isActEnabled = getReactActEnvironment();
+  // TODO
+  const isPortableStory = storyContext.globals.isPortableStory;
 
-  const content = isActEnabled ? (
+  const content = isPortableStory ? (
     <Story {...storyContext} />
   ) : (
     <ErrorBoundary showMain={showMain} showException={showException}>
@@ -80,7 +81,13 @@ export async function renderToCanvas(
     unmountElement(canvasElement);
   }
 
-  await renderElement(element, canvasElement, storyContext?.parameters?.react?.rootOptions);
+  await act(async () => {
+    await renderElement(element, canvasElement, storyContext?.parameters?.react?.rootOptions);
+  });
 
-  return () => unmountElement(canvasElement);
+  return async () => {
+    await act(() => {
+      unmountElement(canvasElement);
+    });
+  };
 }