From b4a1714b725e9e6e87db1f4968fc3f4cbb8c58e2 Mon Sep 17 00:00:00 2001
From: Valentin Palkovic <valentin@chromatic.com>
Date: Thu, 12 Dec 2024 10:22:04 +0100
Subject: [PATCH] React: Use Act wrapper in Storybook for component rendering

---
 code/renderers/react/src/act-compat.ts        |  8 ++++----
 code/renderers/react/src/portable-stories.tsx | 11 +++--------
 code/renderers/react/src/renderToCanvas.tsx   | 17 ++++++++++++-----
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/code/renderers/react/src/act-compat.ts b/code/renderers/react/src/act-compat.ts
index 36e56712e02b..f54a597ca6c8 100644
--- 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
index 7b906c9f4bde..8a0015ddf325 100644
--- 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
index 3ae6136f9582..dcf908bbcff5 100644
--- 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);
+    });
+  };
 }