From b52cbfdb16a103fe3a9f89040fc63a13ce3dded3 Mon Sep 17 00:00:00 2001
From: Max Schmitt <max@schmitt.mx>
Date: Tue, 18 May 2021 18:07:45 +0200
Subject: [PATCH] fix(chromium): close background pages on close (#6608)

---
 src/server/browserContext.ts     | 2 ++
 src/server/chromium/crBrowser.ts | 9 +++++++++
 src/server/firefox/ffBrowser.ts  | 2 ++
 src/server/webkit/wkBrowser.ts   | 2 ++
 tests/chromium/launcher.spec.ts  | 3 +--
 5 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/server/browserContext.ts b/src/server/browserContext.ts
index d8b40ff055..0849dfd968 100644
--- a/src/server/browserContext.ts
+++ b/src/server/browserContext.ts
@@ -151,6 +151,7 @@ export abstract class BrowserContext extends SdkObject {
   abstract _doExposeBinding(binding: PageBinding): Promise<void>;
   abstract _doUpdateRequestInterception(): Promise<void>;
   abstract _doClose(): Promise<void>;
+  abstract _onClosePersistent(): Promise<void>;
 
   async cookies(urls: string | string[] | undefined = []): Promise<types.NetworkCookie[]> {
     if (urls && !Array.isArray(urls))
@@ -283,6 +284,7 @@ export abstract class BrowserContext extends SdkObject {
         // Close all the pages instead of the context,
         // because we cannot close the default context.
         await Promise.all(this.pages().map(page => page.close(metadata)));
+        await this._onClosePersistent();
       } else {
         // Close the context.
         await this._doClose();
diff --git a/src/server/chromium/crBrowser.ts b/src/server/chromium/crBrowser.ts
index 877dddd80e..3748f5f7c1 100644
--- a/src/server/chromium/crBrowser.ts
+++ b/src/server/chromium/crBrowser.ts
@@ -437,6 +437,15 @@ export class CRBrowserContext extends BrowserContext {
     }
   }
 
+  async _onClosePersistent() {
+    for (const [targetId, backgroundPage] of this._browser._backgroundPages.entries()) {
+      if (backgroundPage._browserContext === this && backgroundPage._initializedPage) {
+        backgroundPage.didClose();
+        this._browser._backgroundPages.delete(targetId);
+      }
+    }
+  }
+
   backgroundPages(): Page[] {
     const result: Page[] = [];
     for (const backgroundPage of this._browser._backgroundPages.values()) {
diff --git a/src/server/firefox/ffBrowser.ts b/src/server/firefox/ffBrowser.ts
index 0ad5648425..bf12f0f644 100644
--- a/src/server/firefox/ffBrowser.ts
+++ b/src/server/firefox/ffBrowser.ts
@@ -311,6 +311,8 @@ export class FFBrowserContext extends BrowserContext {
     await this._browser._connection.send('Browser.setRequestInterception', { browserContextId: this._browserContextId, enabled: !!this._requestInterceptor });
   }
 
+  async _onClosePersistent() {}
+
   async _doClose() {
     assert(this._browserContextId);
     await this._browser._connection.send('Browser.removeBrowserContext', { browserContextId: this._browserContextId });
diff --git a/src/server/webkit/wkBrowser.ts b/src/server/webkit/wkBrowser.ts
index b3a9b46361..1a492bf8e0 100644
--- a/src/server/webkit/wkBrowser.ts
+++ b/src/server/webkit/wkBrowser.ts
@@ -317,6 +317,8 @@ export class WKBrowserContext extends BrowserContext {
       await (page._delegate as WKPage).updateRequestInterception();
   }
 
+  async _onClosePersistent() {}
+
   async _doClose() {
     assert(this._browserContextId);
     await this._browser._browserSession.send('Playwright.deleteContext', { browserContextId: this._browserContextId });
diff --git a/tests/chromium/launcher.spec.ts b/tests/chromium/launcher.spec.ts
index 5089ade735..dcd13e67a8 100644
--- a/tests/chromium/launcher.spec.ts
+++ b/tests/chromium/launcher.spec.ts
@@ -72,8 +72,7 @@ it('should return background pages', async ({browserType, browserOptions, create
   expect(context.pages()).not.toContain(backgroundPage);
   await context.close();
   expect(context.pages().length).toBe(0);
-  // TODO: the following line is flaky, uncomment once fixed.
-  // expect(context.backgroundPages().length).toBe(0);
+  expect(context.backgroundPages().length).toBe(0);
 });
 
 it('should return background pages when recording video', async ({browserType, browserOptions, createUserDataDir, asset}, testInfo) => {