From 632402769c40be7c979a0095e16ee5c30667900b Mon Sep 17 00:00:00 2001
From: Yann Braga <yannbf@gmail.com>
Date: Fri, 13 Dec 2024 13:12:20 +0100
Subject: [PATCH 1/3] Build: Improve waiting detection on E2E to fix flake

---
 code/e2e-tests/util.ts | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/code/e2e-tests/util.ts b/code/e2e-tests/util.ts
index 057dc91f715d..0e8c6d0b93d8 100644
--- a/code/e2e-tests/util.ts
+++ b/code/e2e-tests/util.ts
@@ -112,6 +112,10 @@ export class SbPage {
     const storyLoadingPage = root.locator('.sb-preparing-story');
     await docsLoadingPage.waitFor({ state: 'hidden' });
     await storyLoadingPage.waitFor({ state: 'hidden' });
+    await this.page.waitForFunction(() => {
+      const storybookRoot = document.querySelector('#storybook-root');
+      return storybookRoot && storybookRoot.children.length > 0;
+    });
   }
 
   previewIframe() {

From cf809b93e46286e5a446153e6d567bf1428d52bd Mon Sep 17 00:00:00 2001
From: Yann Braga <yannbf@gmail.com>
Date: Fri, 13 Dec 2024 13:15:51 +0100
Subject: [PATCH 2/3] add deps optimization in react kitchen sink

---
 .../react/.storybook/main.ts                     | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/test-storybooks/portable-stories-kitchen-sink/react/.storybook/main.ts b/test-storybooks/portable-stories-kitchen-sink/react/.storybook/main.ts
index 0a76f80b4d25..57b96d2e8d1e 100644
--- a/test-storybooks/portable-stories-kitchen-sink/react/.storybook/main.ts
+++ b/test-storybooks/portable-stories-kitchen-sink/react/.storybook/main.ts
@@ -12,9 +12,21 @@ const config: StorybookConfig = {
     options: {},
   },
   core: {
-    disableWhatsNewNotifications: true
+    disableWhatsNewNotifications: true,
   },
-  previewHead: (head = '') => `${head}
+  viteFinal: (config) => ({
+    ...config,
+    optimizeDeps: {
+      ...config.optimizeDeps,
+      include: [
+        ...(config.optimizeDeps?.include || []),
+        "react-dom/test-utils",
+        "@storybook/react/**",
+        "@storybook/experimental-addon-test/preview",
+      ],
+    },
+  }),
+  previewHead: (head = "") => `${head}
   <style>
     body {
       border: 1px solid red;

From a30a5f39875ff9f2f5d0ec9a462e3b27fdaf843c Mon Sep 17 00:00:00 2001
From: Yann Braga <yannbf@gmail.com>
Date: Fri, 13 Dec 2024 19:43:05 +0100
Subject: [PATCH 3/3] improve story loading

---
 code/e2e-tests/util.ts | 52 ++++++++++++++++++++++++++++++++++++++----
 1 file changed, 47 insertions(+), 5 deletions(-)

diff --git a/code/e2e-tests/util.ts b/code/e2e-tests/util.ts
index 0e8c6d0b93d8..d9faad22939a 100644
--- a/code/e2e-tests/util.ts
+++ b/code/e2e-tests/util.ts
@@ -1,3 +1,4 @@
+/* eslint-disable local-rules/no-uncategorized-errors */
 import { toId } from '@storybook/csf';
 
 import type { Expect, Page } from '@playwright/test';
@@ -61,6 +62,7 @@ export class SbPage {
     await this.expect(selected).toHaveAttribute('data-selected', 'true');
 
     await this.previewRoot();
+    await this.waitUntilLoaded();
   }
 
   async navigateToUnattachedDocs(title: string, name = 'docs') {
@@ -80,7 +82,25 @@ export class SbPage {
     const selected = storyLink;
     await this.expect(selected).toHaveAttribute('data-selected', 'true');
 
-    await this.previewRoot();
+    await this.waitForStoryLoaded();
+  }
+
+  async waitForStoryLoaded() {
+    try {
+      const root = this.previewRoot();
+      // Wait until there is at least one child (a story element) in the preview iframe
+      await root.locator(':scope > *').first().waitFor({
+        state: 'attached',
+        timeout: 10000,
+      });
+    } catch (error: any) {
+      if (error.name === 'TimeoutError') {
+        throw new Error(
+          'The Storybook iframe did not have children within the specified timeout. Did the story load correctly?'
+        );
+      }
+      throw error;
+    }
   }
 
   async waitUntilLoaded() {
@@ -112,10 +132,8 @@ export class SbPage {
     const storyLoadingPage = root.locator('.sb-preparing-story');
     await docsLoadingPage.waitFor({ state: 'hidden' });
     await storyLoadingPage.waitFor({ state: 'hidden' });
-    await this.page.waitForFunction(() => {
-      const storybookRoot = document.querySelector('#storybook-root');
-      return storybookRoot && storybookRoot.children.length > 0;
-    });
+
+    await this.waitForStoryLoaded();
   }
 
   previewIframe() {
@@ -147,6 +165,30 @@ export class SbPage {
   getCanvasBodyElement() {
     return this.previewIframe().locator('body');
   }
+
+  // utility to try and decrease flake
+  async retryTimes(
+    fn: () => Promise<void>,
+    options?: {
+      retries?: number;
+      delay?: number;
+    }
+  ): Promise<void> {
+    let attempts = 0;
+    const { retries = 3, delay = 0 } = options || {};
+    while (attempts < retries) {
+      try {
+        await fn();
+        return;
+      } catch (error) {
+        attempts++;
+        if (attempts === retries) {
+          throw error;
+        }
+        await new Promise<void>((resolve) => setTimeout(resolve, delay));
+      }
+    }
+  }
 }
 
 const templateName: keyof typeof allTemplates = process.env.STORYBOOK_TEMPLATE_NAME || ('' as any);