From 0d9702e75bc7bc60debf683856c65c9ef98243cd Mon Sep 17 00:00:00 2001
From: emiryy <42748013+emiryy@users.noreply.github.com>
Date: Fri, 19 Aug 2022 14:03:28 -0400
Subject: [PATCH] change waitUntil param to take options

changeset
---
 .changeset/lemon-dragons-count.md         | 5 +++++
 packages/storybook-a11y-test/src/index.ts | 9 ++++++---
 2 files changed, 11 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/lemon-dragons-count.md

diff --git a/.changeset/lemon-dragons-count.md b/.changeset/lemon-dragons-count.md
new file mode 100644
index 0000000000..08a22f8f31
--- /dev/null
+++ b/.changeset/lemon-dragons-count.md
@@ -0,0 +1,5 @@
+---
+'@shopify/storybook-a11y-test': minor
+---
+
+Allow waitUntil param to take options
diff --git a/packages/storybook-a11y-test/src/index.ts b/packages/storybook-a11y-test/src/index.ts
index 0817e8901d..254395d4cf 100644
--- a/packages/storybook-a11y-test/src/index.ts
+++ b/packages/storybook-a11y-test/src/index.ts
@@ -2,7 +2,7 @@
 import os from 'os';
 import fs from 'fs';
 
-import puppeteer, {Browser} from 'puppeteer';
+import puppeteer, {Browser, PuppeteerLifeCycleEvent} from 'puppeteer';
 import pMap from 'p-map';
 import chalk from 'chalk';
 import Koa from 'koa';
@@ -57,11 +57,13 @@ export class A11yTestRunner {
     storyIds = [],
     concurrentCount = os.cpus().length,
     timeout = 3000,
+    waitUntil = 'load',
     disableAnimation = false,
   }: {
     storyIds: StoryId[];
     concurrentCount?: number;
     timeout?: number;
+    waitUntil?: PuppeteerLifeCycleEvent;
     disableAnimation?: boolean;
   }) {
     console.log(
@@ -75,7 +77,7 @@ export class A11yTestRunner {
     console.log(chalk.bold(`🧪 Testing ${storyIds.length} urls with axe`));
     const results = await pMap(
       storyIds,
-      this.generateTestStoryFunction(timeout, disableAnimation),
+      this.generateTestStoryFunction(timeout, waitUntil, disableAnimation),
       {
         concurrency: concurrentCount,
       },
@@ -146,6 +148,7 @@ export class A11yTestRunner {
 
   private generateTestStoryFunction(
     timeout: number,
+    waitUntil: PuppeteerLifeCycleEvent,
     disableAnimation: boolean,
   ) {
     return async (id: StoryId) => {
@@ -154,7 +157,7 @@ export class A11yTestRunner {
 
       try {
         await page.goto(`${this.iframePath}?id=${id}`, {
-          waitUntil: 'load',
+          waitUntil,
           timeout,
         });