From cf69af5c7f0698cdd57c22a2af795f036f2b274e Mon Sep 17 00:00:00 2001
From: Rui Figueira <rui.figueira@gmail.com>
Date: Wed, 8 Jan 2025 00:04:34 +0000
Subject: [PATCH] fix: closing default crxApp did not reset crxTransport

See: https://github.com/ruifigueira/playwright-crx/issues/7#issuecomment-2570047648
---
 src/server/crx.ts     | 24 ++++++++++++------------
 tests/crx/api.spec.ts |  7 ++++++-
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/server/crx.ts b/src/server/crx.ts
index 5023314b7..5170c8481 100644
--- a/src/server/crx.ts
+++ b/src/server/crx.ts
@@ -49,8 +49,8 @@ export function tabIdFromPage(page: Page): number | undefined {
 
 export class Crx extends SdkObject {
 
-  private _transport!: CrxTransport;
-  private _browserPromise!: Promise<CRBrowser>;
+  private _transport?: CrxTransport;
+  private _browserPromise?: Promise<CRBrowser>;
   private _crxApplicationPromise: Promise<CrxApplication> | undefined;
   private _incognitoCrxApplicationPromise: Promise<CrxApplication> | undefined;
   readonly player: CrxPlayer;
@@ -77,7 +77,7 @@ export class Crx extends SdkObject {
         process: undefined,
         // browser.close() calls this function, and closing transport will trigger
         // Browser.Events.Disconnected and force the browser to resolve the close promise.
-        close: () => this._transport.closeAndWait(),
+        close: () => this._transport!.closeAndWait(),
         kill: () => Promise.resolve(),
       };
       const browserOptions: BrowserOptions = {
@@ -97,17 +97,18 @@ export class Crx extends SdkObject {
       this._transport = new CrxTransport();
       this._browserPromise = CRBrowser.connect(this.attribution.playwright, this._transport, browserOptions);
     }
-    const browser = await this._browserPromise;
+    const browser = await this._browserPromise!;
+    const transport = this._transport!;
 
     if (incognito) {
       if (this._incognitoCrxApplicationPromise)
         throw new Error(`incognito crxApplication is already started`);
-      this._incognitoCrxApplicationPromise = this._startIncognitoCrxApplication(browser, this._transport, newContextOptions);
+      this._incognitoCrxApplicationPromise = this._startIncognitoCrxApplication(browser, transport, newContextOptions);
       return await this._incognitoCrxApplicationPromise;
     } else {
       if (this._crxApplicationPromise)
         throw new Error(`crxApplication is already started`);
-      this._crxApplicationPromise = this._startCrxApplication(browser, this._transport);
+      this._crxApplicationPromise = this._startCrxApplication(browser, transport);
       return await this._crxApplicationPromise;
     }
   }
@@ -118,6 +119,10 @@ export class Crx extends SdkObject {
     context.on(BrowserContext.Events.Close, () => {
       this._crxApplicationPromise = undefined;
     });
+    browser.on(CRBrowser.Events.Disconnected, () => {
+      this._browserPromise = undefined;
+      this._transport = undefined;
+    });
     // override factory otherwise it will fail because the default factory tries to launch a new playwright app
     RecorderApp.factory = (): IRecorderAppFactory => {
       return async recorder => {
@@ -159,7 +164,7 @@ export class Crx extends SdkObject {
     context.on(BrowserContext.Events.Close, () => {
       this._incognitoCrxApplicationPromise = undefined;
     });
-    const crxApp = new CrxApplication(this, context, this._transport);
+    const crxApp = new CrxApplication(this, context, transport);
     await crxApp.attach(incognitoTabId);
     return crxApp;
   }
@@ -169,11 +174,6 @@ export class Crx extends SdkObject {
       await this._incognitoCrxApplicationPromise :
       await this._crxApplicationPromise;
   }
-
-  async closeAndWait() {
-    if (this._browserPromise)
-      await this._browserPromise.then(browser => browser.close({}));
-  }
 }
 
 export class CrxApplication extends SdkObject {
diff --git a/tests/crx/api.spec.ts b/tests/crx/api.spec.ts
index c00bec195..78d4e71cc 100644
--- a/tests/crx/api.spec.ts
+++ b/tests/crx/api.spec.ts
@@ -200,7 +200,7 @@ test('should close crxApp', async ({ runCrxTest }) => {
   await runCrxTest(async ({ crx, crxApp, expect }) => {
     await crxApp.close();
     await expect(crxApp.newPage()).rejects.toThrow();
-    await expect(await crx.get()).toBeUndefined();
+    expect(await crx.get()).toBeUndefined();
   });
 });
 
@@ -210,5 +210,10 @@ test('should be able to open new crxApp after close', async ({ runCrxTest }) =>
     expect(await crx.get()).toBeUndefined();
     const newCrxApp = await crx.start();
     expect(await crx.get()).toBe(newCrxApp);
+
+    // ensure we can open a new page on the newly opened crxApp
+    // see: https://github.com/ruifigueira/playwright-crx/issues/7#issuecomment-2570047648
+    const newPage = await newCrxApp.newPage();
+    expect(newPage).toBeTruthy();
   });
 });