From 62839d8e81c1d68bbc177ff35c4f1dd03cad8496 Mon Sep 17 00:00:00 2001
From: "Manu Mtz.-Almeida" <manu.mtza@gmail.com>
Date: Thu, 20 Jun 2019 18:47:36 +0200
Subject: [PATCH] feat(testing): expose waitFor option

---
 .../puppeteer/puppeteer-declarations.ts       |  6 +++---
 src/testing/puppeteer/puppeteer-page.ts       | 20 +++++++++++++------
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/testing/puppeteer/puppeteer-declarations.ts b/src/testing/puppeteer/puppeteer-declarations.ts
index c965c1b1776..a436cf1c66a 100644
--- a/src/testing/puppeteer/puppeteer-declarations.ts
+++ b/src/testing/puppeteer/puppeteer-declarations.ts
@@ -2,7 +2,7 @@ import * as d from '../../declarations';
 import * as puppeteer from 'puppeteer';
 
 
-export interface NewE2EPageOptions {
+export interface NewE2EPageOptions extends puppeteer.NavigationOptions {
   url?: string;
   html?: string;
 }
@@ -80,7 +80,7 @@ export interface E2EPage extends PuppeteerPage {
    * a localhost address. A shortcut to `page.goto(url)` is to set the `url` option
    * when creating a new page, such as `const page = await newE2EPage({ url })`.
    */
-  goTo(url: string, options?: Partial<puppeteer.NavigationOptions>): Promise<puppeteer.Response | null>;
+  goTo(url: string, options?: puppeteer.NavigationOptions): Promise<puppeteer.Response | null>;
 
   /**
    * Instead of testing a url directly, html content can be mocked using
@@ -88,7 +88,7 @@ export interface E2EPage extends PuppeteerPage {
    * the `html` option when creating a new page, such as
    * `const page = await newE2EPage({ html })`.
    */
-  setContent(html: string): Promise<void>;
+  setContent(html: string, options?: puppeteer.NavigationOptions): Promise<void>;
 
   /**
    * Used to test if an event was, or was not dispatched. This method
diff --git a/src/testing/puppeteer/puppeteer-page.ts b/src/testing/puppeteer/puppeteer-page.ts
index f8b22596cab..686f084ba46 100644
--- a/src/testing/puppeteer/puppeteer-page.ts
+++ b/src/testing/puppeteer/puppeteer-page.ts
@@ -58,13 +58,13 @@ export async function newE2EPage(opts: pd.NewE2EPageOptions = {}): Promise<pd.E2
     page.on('requestfailed', requestFailed);
 
     if (typeof opts.html === 'string') {
-      const errMsg = await e2eSetContent(page, opts.html);
+      const errMsg = await e2eSetContent(page, opts.html, { waitUntil: opts.waitUntil });
       if (errMsg) {
         throw errMsg;
       }
 
     } else if (typeof opts.url === 'string') {
-      const errMsg = await e2eGoTo(page, opts.url);
+      const errMsg = await e2eGoTo(page, opts.url, { waitUntil: opts.waitUntil });
       if (errMsg) {
         throw errMsg;
       }
@@ -93,7 +93,7 @@ function failedPage() {
   return page;
 }
 
-async function e2eGoTo(page: pd.E2EPageInternal, url: string) {
+async function e2eGoTo(page: pd.E2EPageInternal, url: string, options: puppeteer.NavigationOptions) {
   try {
     if (page.isClosed()) {
       return 'e2eGoTo unavailable: page already closed';
@@ -120,7 +120,10 @@ async function e2eGoTo(page: pd.E2EPageInternal, url: string) {
 
   const fullUrl = browserUrl + url.substring(1);
 
-  const rsp = await page._e2eGoto(fullUrl, { waitUntil: 'networkidle0' });
+  if (!options.waitUntil) {
+    options.waitUntil = DEFAULT_WAIT_FOR;
+  }
+  const rsp = await page._e2eGoto(fullUrl, options);
 
   if (!rsp.ok()) {
     await closePage(page);
@@ -140,7 +143,7 @@ async function e2eGoTo(page: pd.E2EPageInternal, url: string) {
 }
 
 
-async function e2eSetContent(page: pd.E2EPageInternal, html: string) {
+async function e2eSetContent(page: pd.E2EPageInternal, html: string, options: puppeteer.NavigationOptions = {}) {
   try {
     if (page.isClosed()) {
       return 'e2eSetContent unavailable: page already closed';
@@ -180,7 +183,10 @@ async function e2eSetContent(page: pd.E2EPageInternal, html: string) {
     }
   });
 
-  const rsp = await page._e2eGoto(url, { waitUntil: 'networkidle0' });
+  if (!options.waitUntil) {
+    options.waitUntil = DEFAULT_WAIT_FOR;
+  }
+  const rsp = await page._e2eGoto(url, options);
 
   if (!rsp.ok()) {
     await closePage(page);
@@ -310,3 +316,5 @@ function pageError(e: any) {
 function requestFailed(req: puppeteer.Request) {
   console.error('requestfailed', req.url());
 }
+
+const DEFAULT_WAIT_FOR = 'load';