diff --git a/code/lib/store/template/stories/rendering.stories.ts b/code/lib/store/template/stories/rendering.stories.ts
index 4c00bf8e32dc..f7388fe59d61 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 54260e7825aa..9ec2fffd5264 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));
     }