diff --git a/src/common/DOMWorld.ts b/src/common/DOMWorld.ts
index 14d23d0418f13..9439c2fbe3de0 100644
--- a/src/common/DOMWorld.ts
+++ b/src/common/DOMWorld.ts
@@ -645,6 +645,7 @@ export class DOMWorld {
     const waitTaskOptions: WaitTaskOptions = {
       domWorld: this,
       predicateBody: helper.makePredicateString(predicate, queryOne),
+      predicateAcceptsContextElement: true,
       title,
       polling,
       timeout,
@@ -691,6 +692,7 @@ export class DOMWorld {
     const waitTaskOptions: WaitTaskOptions = {
       domWorld: this,
       predicateBody: helper.makePredicateString(predicate),
+      predicateAcceptsContextElement: true,
       title,
       polling,
       timeout,
@@ -717,6 +719,7 @@ export class DOMWorld {
     const waitTaskOptions: WaitTaskOptions = {
       domWorld: this,
       predicateBody: pageFunction,
+      predicateAcceptsContextElement: false,
       title: 'function',
       polling,
       timeout,
@@ -737,6 +740,7 @@ export class DOMWorld {
 export interface WaitTaskOptions {
   domWorld: DOMWorld;
   predicateBody: Function | string;
+  predicateAcceptsContextElement: boolean;
   title: string;
   polling: string | number;
   timeout: number;
@@ -753,6 +757,7 @@ export class WaitTask {
   _polling: string | number;
   _timeout: number;
   _predicateBody: string;
+  _predicateAcceptsContextElement: boolean;
   _args: SerializableOrJSHandle[];
   _binding: PageBinding;
   _runCount = 0;
@@ -786,6 +791,8 @@ export class WaitTask {
     this._timeout = options.timeout;
     this._root = options.root;
     this._predicateBody = getPredicateBody(options.predicateBody);
+    this._predicateAcceptsContextElement =
+      options.predicateAcceptsContextElement;
     this._args = options.args;
     this._binding = options.binding;
     this._runCount = 0;
@@ -836,6 +843,7 @@ export class WaitTask {
           waitForPredicatePageFunction,
           this._root,
           this._predicateBody,
+          this._predicateAcceptsContextElement,
           this._polling,
           this._timeout,
           ...this._args
@@ -845,6 +853,7 @@ export class WaitTask {
           waitForPredicatePageFunction,
           null,
           this._predicateBody,
+          this._predicateAcceptsContextElement,
           this._polling,
           this._timeout,
           ...this._args
@@ -912,6 +921,7 @@ export class WaitTask {
 async function waitForPredicatePageFunction(
   root: Element | Document | null,
   predicateBody: string,
+  predicateAcceptsContextElement: boolean,
   polling: string,
   timeout: number,
   ...args: unknown[]
@@ -928,7 +938,9 @@ async function waitForPredicatePageFunction(
    * @returns {!Promise<*>}
    */
   async function pollMutation(): Promise<unknown> {
-    const success = await predicate(root, ...args);
+    const success = predicateAcceptsContextElement
+      ? await predicate(root, ...args)
+      : await predicate(...args);
     if (success) return Promise.resolve(success);
 
     let fulfill;
@@ -938,7 +950,9 @@ async function waitForPredicatePageFunction(
         observer.disconnect();
         fulfill();
       }
-      const success = await predicate(root, ...args);
+      const success = predicateAcceptsContextElement
+        ? await predicate(root, ...args)
+        : await predicate(...args);
       if (success) {
         observer.disconnect();
         fulfill(success);
@@ -963,7 +977,9 @@ async function waitForPredicatePageFunction(
         fulfill();
         return;
       }
-      const success = await predicate(root, ...args);
+      const success = predicateAcceptsContextElement
+        ? await predicate(root, ...args)
+        : await predicate(...args);
       if (success) fulfill(success);
       else requestAnimationFrame(onRaf);
     }
@@ -980,7 +996,9 @@ async function waitForPredicatePageFunction(
         fulfill();
         return;
       }
-      const success = await predicate(root, ...args);
+      const success = predicateAcceptsContextElement
+        ? await predicate(root, ...args)
+        : await predicate(...args);
       if (success) fulfill(success);
       else setTimeout(onTimeout, pollInterval);
     }
diff --git a/test/network.spec.ts b/test/network.spec.ts
index e09633983f619..4e098b196400d 100644
--- a/test/network.spec.ts
+++ b/test/network.spec.ts
@@ -592,7 +592,7 @@ describe('network', function () {
       expect(requests.get('script.js').isNavigationRequest()).toBe(false);
       expect(requests.get('style.css').isNavigationRequest()).toBe(false);
     });
-    it('should work when navigating to image', async () => {
+    itFailsFirefox('should work when navigating to image', async () => {
       const { page, server } = getTestState();
 
       const requests = [];
diff --git a/test/waittask.spec.ts b/test/waittask.spec.ts
index bce546adbb4e1..7b7f409e8cea1 100644
--- a/test/waittask.spec.ts
+++ b/test/waittask.spec.ts
@@ -293,7 +293,11 @@ describe('waittask specs', function () {
       const div = await page.$('div');
       let resolved = false;
       const waitForFunction = page
-        .waitForFunction((element) => !element.parentElement, {}, div)
+        .waitForFunction(
+          (element) => element.localName === 'div' && !element.parentElement,
+          {},
+          div
+        )
         .then(() => (resolved = true));
       expect(resolved).toBe(false);
       await page.evaluate((element: HTMLElement) => element.remove(), div);