fix(route): support route w/ async handler & times (#14317)

microsoft · May 22, 2022 · a1324bd · a1324bd
1 parent abed166
commit a1324bd
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 23 deletions.
diff --git a/packages/playwright-core/src/client/browserContext.ts b/packages/playwright-core/src/client/browserContext.ts
@@ -145,19 +145,26 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel>
 
   _onRoute(route: network.Route, request: network.Request) {
     for (const routeHandler of this._routes) {
-      if (routeHandler.matches(request.url())) {
+      if (!routeHandler.matches(request.url()))
+        continue;
+      // Immediately deactivate based on |times|.
+      if (routeHandler.willExpire())
+        this._routes.splice(this._routes.indexOf(routeHandler), 1);
+
+      (async () => {
         try {
-          routeHandler.handle(route, request);
+          // Let async callback work prior to disabling interception.
+          await routeHandler.handle(route, request);
         } finally {
-          if (!routeHandler.isActive()) {
-            this._routes.splice(this._routes.indexOf(routeHandler), 1);
-            if (!this._routes.length)
-              this._wrapApiCall(() => this._disableInterception(), true).catch(() => {});
-          }
+          if (!this._routes.length)
+            await this._wrapApiCall(() => this._disableInterception(), true).catch(() => {});
         }
-        return;
-      }
+      })();
+
+      // There is no chaining, first handler wins.
+      return;
     }
+
     // it can race with BrowserContext.close() which then throws since its closed
     route._internalContinue();
   }

diff --git a/packages/playwright-core/src/client/network.ts b/packages/playwright-core/src/client/network.ts
@@ -309,7 +309,7 @@ export class Route extends ChannelOwner<channels.RouteChannel> implements api.Ro
   }
 }
 
-export type RouteHandlerCallback = (route: Route, request: Request) => void;
+export type RouteHandlerCallback = (route: Route, request: Request) => void | Promise<void>;
 
 export type ResourceTiming = {
   startTime: number;
@@ -518,13 +518,13 @@ export class RouteHandler {
     return urlMatches(this._baseURL, requestURL, this.url);
   }
 
-  public handle(route: Route, request: Request): void {
+  public handle(route: Route, request: Request): Promise<void> | void {
     ++this.handledCount;
-    this.handler(route, request);
+    return this.handler(route, request);
   }
 
-  public isActive(): boolean {
-    return this.handledCount < this._times;
+  public willExpire(): boolean {
+    return this.handledCount + 1 >= this._times;
   }
 }
 

diff --git a/packages/playwright-core/src/client/page.ts b/packages/playwright-core/src/client/page.ts
@@ -180,18 +180,24 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
 
   private _onRoute(route: Route, request: Request) {
     for (const routeHandler of this._routes) {
-      if (routeHandler.matches(request.url())) {
+      if (!routeHandler.matches(request.url()))
+        continue;
+      // Immediately deactivate based on |times|.
+      if (routeHandler.willExpire())
+        this._routes.splice(this._routes.indexOf(routeHandler), 1);
+
+      (async () => {
         try {
-          routeHandler.handle(route, request);
+          // Let async callback work prior to disabling interception.
+          await routeHandler.handle(route, request);
         } finally {
-          if (!routeHandler.isActive()) {
-            this._routes.splice(this._routes.indexOf(routeHandler), 1);
-            if (!this._routes.length)
-              this._wrapApiCall(() => this._disableInterception(), true).catch(() => {});
-          }
+          if (!this._routes.length)
+            this._wrapApiCall(() => this._disableInterception(), true).catch(() => {});
         }
-        return;
-      }
+      })();
+
+      // There is no chaining, first handler wins.
+      return;
     }
     this._browserContext._onRoute(route, request);
   }

diff --git a/tests/library/browsercontext-route.spec.ts b/tests/library/browsercontext-route.spec.ts
@@ -211,6 +211,20 @@ it('should support the times parameter with route matching', async ({ context, p
   expect(intercepted).toHaveLength(1);
 });
 
+it('should support async handler w/ times', async ({ context, page, server }) => {
+  await context.route('**/empty.html', async route => {
+    await new Promise(f => setTimeout(f, 100));
+    route.fulfill({
+      body: '<html>intercepted</html>',
+      contentType: 'text/html'
+    });
+  }, { times: 1 });
+  await page.goto(server.EMPTY_PAGE);
+  await expect(page.locator('body')).toHaveText('intercepted');
+  await page.goto(server.EMPTY_PAGE);
+  await expect(page.locator('body')).not.toHaveText('intercepted');
+});
+
 it('should overwrite post body with empty string', async ({ context, server, page, browserName }) => {
   await context.route('**/empty.html', route => {
     route.continue({

diff --git a/tests/page/page-route.spec.ts b/tests/page/page-route.spec.ts
@@ -772,6 +772,20 @@ it('should support the times parameter with route matching', async ({ page, serv
   expect(intercepted).toHaveLength(1);
 });
 
+it('should support async handler w/ times', async ({ page, server }) => {
+  await page.route('**/empty.html', async route => {
+    await new Promise(f => setTimeout(f, 100));
+    route.fulfill({
+      body: '<html>intercepted</html>',
+      contentType: 'text/html'
+    });
+  }, { times: 1 });
+  await page.goto(server.EMPTY_PAGE);
+  await expect(page.locator('body')).toHaveText('intercepted');
+  await page.goto(server.EMPTY_PAGE);
+  await expect(page.locator('body')).not.toHaveText('intercepted');
+});
+
 it('should contain raw request header', async ({ page, server }) => {
   let headers: any;
   await page.route('**/*', async route => {