From 0ad2ef76f66b6a5acd6c06f1355b9a5af742a4db Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@chromatic.com>
Date: Wed, 7 Sep 2022 20:41:55 +1000
Subject: [PATCH 1/2] Fix issue with react 18 implementation

---
 .../store/template/stories/rendering.stories.ts    | 14 +++++++-------
 code/renderers/react/src/render.tsx                | 13 +++++--------
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/code/lib/store/template/stories/rendering.stories.ts b/code/lib/store/template/stories/rendering.stories.ts
index 61519942b63a..c8e069e893fb 100644
--- a/code/lib/store/template/stories/rendering.stories.ts
+++ b/code/lib/store/template/stories/rendering.stories.ts
@@ -30,13 +30,13 @@ export const ChangeArgs = {
     await button.focus();
     await expect(button).toHaveFocus();
 
-    // // When we change the args to the button, it should not rerender
-    // await channel.emit('updateStoryArgs', { storyId: id, updatedArgs: { children: 'New Text' } });
-    // await within(canvasElement).findByText(/New Text/);
-    // await expect(button).toHaveFocus();
+    // When we change the args to the button, it should not rerender
+    await channel.emit('updateStoryArgs', { storyId: id, updatedArgs: { children: 'New Text' } });
+    await within(canvasElement).findByText(/New Text/);
+    await expect(button).toHaveFocus();
 
-    // await channel.emit('resetStoryArgs', { storyId: id });
-    // await within(canvasElement).findByText(/Click me/);
-    // await expect(button).toHaveFocus();
+    await channel.emit('resetStoryArgs', { storyId: id });
+    await within(canvasElement).findByText(/Click me/);
+    await expect(button).toHaveFocus();
   },
 };
diff --git a/code/renderers/react/src/render.tsx b/code/renderers/react/src/render.tsx
index df56d64da981..2fe366999e46 100644
--- a/code/renderers/react/src/render.tsx
+++ b/code/renderers/react/src/render.tsx
@@ -39,10 +39,10 @@ const WithCallback: FC<{ callback: () => void; children: ReactElement }> = ({
   children,
 }) => {
   // See https://github.com/reactwg/react-18/discussions/5#discussioncomment-2276079
-  const once = useRef(false);
+  const once = useRef<() => void>();
   useLayoutEffect(() => {
-    if (once.current) return;
-    once.current = true;
+    if (once.current === callback) return;
+    once.current = callback;
     callback();
   }, [callback]);
 
@@ -55,11 +55,7 @@ const renderElement = async (node: ReactElement, el: Element) => {
 
   return new Promise((resolve) => {
     if (root) {
-      root.render(
-        <WithCallback key={Math.random()} callback={() => resolve(null)}>
-          {node}
-        </WithCallback>
-      );
+      root.render(<WithCallback callback={() => resolve(null)}>{node}</WithCallback>);
     } else {
       ReactDOM.render(node, el, () => resolve(null));
     }
@@ -160,6 +156,7 @@ export async function renderToDOM(
   // This could leads to issues like below:
   // https://github.com/storybookjs/react-storybook/issues/81
   // (This is not the case when we change args or globals to the story however)
+  console.log({ forceRemount });
   if (forceRemount) {
     unmountElement(domElement);
   }

From bdf0a82b05b2fc86dad0419c40f760ca29c63e64 Mon Sep 17 00:00:00 2001
From: Michael Shilman <shilman@users.noreply.github.com>
Date: Wed, 7 Sep 2022 20:14:59 +0800
Subject: [PATCH 2/2] Update code/renderers/react/src/render.tsx

---
 code/renderers/react/src/render.tsx | 1 -
 1 file changed, 1 deletion(-)

diff --git a/code/renderers/react/src/render.tsx b/code/renderers/react/src/render.tsx
index 2fe366999e46..d1e29c5316db 100644
--- a/code/renderers/react/src/render.tsx
+++ b/code/renderers/react/src/render.tsx
@@ -156,7 +156,6 @@ export async function renderToDOM(
   // This could leads to issues like below:
   // https://github.com/storybookjs/react-storybook/issues/81
   // (This is not the case when we change args or globals to the story however)
-  console.log({ forceRemount });
   if (forceRemount) {
     unmountElement(domElement);
   }