fix(waitTask): remove rerunnable tasks from the context data upon success/failure (#3875)

src/server/frames.ts

@@ -999,12 +999,14 @@ class RerunnableTask {
   private _reject: (reason: Error) => void = () => {};
   private _progress: Progress;
   private _returnByValue: boolean;
+  private _contextData: ContextData;
 
   constructor(data: ContextData, progress: Progress, task: dom.SchedulableTask<any>, returnByValue: boolean) {
     this._task = task;
     this._progress = progress;
     this._returnByValue = returnByValue;
-    data.rerunnableTasks.add(this);
+    this._contextData = data;
+    this._contextData.rerunnableTasks.add(this);
     this.promise = new Promise<any>((resolve, reject) => {
       // The task is either resolved with a value, or rejected with a meaningful evaluation error.
       this._resolve = resolve;
@@ -1021,6 +1023,7 @@ class RerunnableTask {
       const injectedScript = await context.injectedScript();
       const pollHandler = new dom.InjectedScriptPollHandler(this._progress, await this._task(injectedScript));
       const result = this._returnByValue ? await pollHandler.finish() : await pollHandler.finishHandle();
+      this._contextData.rerunnableTasks.delete(this);
       this._resolve(result);
     } catch (e) {
       // When the page is navigated, the promise is rejected.
@@ -1033,6 +1036,7 @@ class RerunnableTask {
       if (e.message.includes('Cannot find context with specified id'))
         return;
 
+      this._contextData.rerunnableTasks.delete(this);
       this._reject(e);
     }
   }

test/wait-for-function.spec.ts

@@ -210,3 +210,57 @@ it('should work with multiline body', async ({page}) => {
 it('should wait for predicate with arguments', async ({page}) => {
   await page.waitForFunction(({arg1, arg2}) => arg1 + arg2 === 3, { arg1: 1, arg2: 2});
 });
+
+it('should not be called after finishing successfully', async ({page, server}) => {
+  await page.goto(server.EMPTY_PAGE);
+
+  const messages = [];
+  page.on('console', msg => {
+    if (msg.text().startsWith('waitForFunction'))
+      messages.push(msg.text());
+  });
+
+  await page.waitForFunction(() => {
+    console.log('waitForFunction1');
+    return true;
+  });
+  await page.reload();
+  await page.waitForFunction(() => {
+    console.log('waitForFunction2');
+    return true;
+  });
+  await page.reload();
+  await page.waitForFunction(() => {
+    console.log('waitForFunction3');
+    return true;
+  });
+
+  expect(messages.join('|')).toBe('waitForFunction1|waitForFunction2|waitForFunction3');
+});
+
+it('should not be called after finishing unsuccessfully', async ({page, server}) => {
+  await page.goto(server.EMPTY_PAGE);
+
+  const messages = [];
+  page.on('console', msg => {
+    if (msg.text().startsWith('waitForFunction'))
+      messages.push(msg.text());
+  });
+
+  await page.waitForFunction(() => {
+    console.log('waitForFunction1');
+    throw new Error('waitForFunction1');
+  }).catch(e => null);
+  await page.reload();
+  await page.waitForFunction(() => {
+    console.log('waitForFunction2');
+    throw new Error('waitForFunction2');
+  }).catch(e => null);
+  await page.reload();
+  await page.waitForFunction(() => {
+    console.log('waitForFunction3');
+    throw new Error('waitForFunction3');
+  }).catch(e => null);
+
+  expect(messages.join('|')).toBe('waitForFunction1|waitForFunction2|waitForFunction3');
+});