From f96341eef4fc7aeb7523d6bac2ebc7b23cdc8c0e Mon Sep 17 00:00:00 2001
From: Jeppe Reinhold <jeppe@reinhold.is>
Date: Wed, 15 May 2024 23:43:31 +0200
Subject: [PATCH] attempt to get the preview page multiple times

---
 scripts/bench/browse.ts |  8 ++++----
 scripts/bench/utils.ts  | 25 ++++++++++++++++++++-----
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/scripts/bench/browse.ts b/scripts/bench/browse.ts
index 43d03756411b..f95edea759c5 100644
--- a/scripts/bench/browse.ts
+++ b/scripts/bench/browse.ts
@@ -37,7 +37,7 @@ async function benchAutodocs(url: string) {
   const browser = await chromium.launch(/* { headless: false } */);
   await browser.newContext();
   const page = await browser.newPage();
-  await page.setDefaultTimeout(40000);
+  page.setDefaultTimeout(40000);
 
   const start = now();
   await page.goto(`${url}?path=/docs/example-button--docs`);
@@ -45,7 +45,7 @@ async function benchAutodocs(url: string) {
   const tasks = [
     async () => {
       const previewPage = await getPreviewPage(page);
-      await previewPage.setDefaultTimeout(40000);
+      previewPage.setDefaultTimeout(40000);
 
       await previewPage.waitForLoadState('load');
       await previewPage.getByText('Primary UI component for user interaction');
@@ -72,7 +72,7 @@ async function benchMDX(url: string) {
   const tasks = [
     async () => {
       const previewPage = await getPreviewPage(page);
-      await previewPage.setDefaultTimeout(40000);
+      previewPage.setDefaultTimeout(40000);
 
       await previewPage.waitForLoadState('load');
       await previewPage.getByText('Configure your project');
@@ -108,7 +108,7 @@ async function benchStory(url: string) {
     },
     async () => {
       const previewPage = await getPreviewPage(page);
-      await previewPage.setDefaultTimeout(40000);
+      previewPage.setDefaultTimeout(40000);
 
       await previewPage.waitForLoadState('load');
       await previewPage.getByText('Button');
diff --git a/scripts/bench/utils.ts b/scripts/bench/utils.ts
index c9c159cb5782..7457cf80e359 100644
--- a/scripts/bench/utils.ts
+++ b/scripts/bench/utils.ts
@@ -33,9 +33,24 @@ export const loadBench = async (options: SaveBenchOptions): Promise<Partial<Benc
 };
 
 export async function getPreviewPage(page: Page) {
-  await page.waitForFunction(() => {
-    return document.querySelector('iframe')?.contentDocument.readyState === 'complete';
-  });
-  const previewPage = await page.frame({ url: /iframe.html/ }).page();
-  return previewPage;
+  /**
+   * Fix flakiness in preview iframe retrieval
+   * Sometimes the iframe is not yet available when we try to access it,
+   * even after waiting for the readyState to be complete.
+   *
+   * This loop will keep trying to access the iframe until it's available.
+   */
+  for (let i = 0; i < 10; i++) {
+    // eslint-disable-next-line @typescript-eslint/no-loop-func
+    await page.waitForFunction(() => {
+      return document.querySelector('iframe')?.contentDocument.readyState === 'complete';
+    });
+
+    const previewPage = page.frame({ url: /iframe.html/ })?.page();
+    if (previewPage) {
+      return previewPage;
+    }
+  }
+
+  throw new Error('The preview iframe was never found');
 }