From f9d9ad25dec032eec7025978e31b78229d69b951 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Wed, 18 Sep 2024 09:34:06 -0700
Subject: [PATCH] feat(locator handler): perform checkpoit during
 `locator.waitFor` (#32683)

Fixes #32255.
---
 packages/playwright-core/src/server/frames.ts |  7 +++++--
 packages/playwright-core/src/server/page.ts   |  2 +-
 tests/page/page-add-locator-handler.spec.ts   | 18 ++++++++++++++++++
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/packages/playwright-core/src/server/frames.ts b/packages/playwright-core/src/server/frames.ts
index 3b952ea02a..b7b626d09f 100644
--- a/packages/playwright-core/src/server/frames.ts
+++ b/packages/playwright-core/src/server/frames.ts
@@ -782,13 +782,16 @@ export class Frame extends SdkObject {
       throw new Error(`state: expected one of (attached|detached|visible|hidden)`);
     return controller.run(async progress => {
       progress.log(`waiting for ${this._asLocator(selector)}${state === 'attached' ? '' : ' to be ' + state}`);
-      return await this.waitForSelectorInternal(progress, selector, options, scope);
+      return await this.waitForSelectorInternal(progress, selector, true, options, scope);
     }, this._page._timeoutSettings.timeout(options));
   }
 
-  async waitForSelectorInternal(progress: Progress, selector: string, options: types.WaitForElementOptions, scope?: dom.ElementHandle): Promise<dom.ElementHandle<Element> | null> {
+  async waitForSelectorInternal(progress: Progress, selector: string, performLocatorHandlersCheckpoint: boolean, options: types.WaitForElementOptions, scope?: dom.ElementHandle): Promise<dom.ElementHandle<Element> | null> {
     const { state = 'visible' } = options;
     const promise = this.retryWithProgressAndTimeouts(progress, [0, 20, 50, 100, 100, 500], async continuePolling => {
+      if (performLocatorHandlersCheckpoint)
+        await this._page.performLocatorHandlersCheckpoint(progress);
+
       const resolved = await this.selectors.resolveInjectedForSelector(selector, options, scope);
       progress.throwIfAborted();
       if (!resolved) {
diff --git a/packages/playwright-core/src/server/page.ts b/packages/playwright-core/src/server/page.ts
index aeaeb0af88..04728b681e 100644
--- a/packages/playwright-core/src/server/page.ts
+++ b/packages/playwright-core/src/server/page.ts
@@ -473,7 +473,7 @@ export class Page extends SdkObject {
           progress.throwIfAborted();
           if (!handler.noWaitAfter) {
             progress.log(`  locator handler has finished, waiting for ${asLocator(this.attribution.playwright.options.sdkLanguage, handler.selector)} to be hidden`);
-            await this.mainFrame().waitForSelectorInternal(progress, handler.selector, { state: 'hidden' });
+            await this.mainFrame().waitForSelectorInternal(progress, handler.selector, false, { state: 'hidden' });
           } else {
             progress.log(`  locator handler has finished`);
           }
diff --git a/tests/page/page-add-locator-handler.spec.ts b/tests/page/page-add-locator-handler.spec.ts
index c99dc1d7c4..f8b1bb6ba7 100644
--- a/tests/page/page-add-locator-handler.spec.ts
+++ b/tests/page/page-add-locator-handler.spec.ts
@@ -181,6 +181,24 @@ test('should work with toBeVisible', async ({ page, server }) => {
   expect(called).toBe(1);
 });
 
+test('should work with locator.waitFor', async ({ page, server }) => {
+  await page.goto(server.PREFIX + '/input/handle-locator.html');
+
+  let called = 0;
+  await page.addLocatorHandler(page.getByText('This interstitial covers the button'), async () => {
+    ++called;
+    await page.locator('#close').click();
+  });
+
+  await page.evaluate(() => {
+    (window as any).clicked = 0;
+    (window as any).setupAnnoyingInterstitial('remove', 1);
+  });
+  await page.locator('#target').waitFor();
+  await expect(page.locator('#interstitial')).not.toBeVisible();
+  expect(called).toBe(1);
+});
+
 test('should work with toHaveScreenshot', async ({ page, server, isAndroid }) => {
   test.fixme(isAndroid, 'Screenshots are cut off on Android');
   await page.setViewportSize({ width: 500, height: 500 });