From a2d8e2eacb803a7cc6113a472fe0ccccf710cf5c Mon Sep 17 00:00:00 2001
From: Roman Kuznetsov <kuznetsov.roman.orenburg@gmail.com>
Date: Tue, 27 Aug 2024 17:19:23 +0300
Subject: [PATCH] fix: add getPuppeteer timeout

---
 src/browser/commands/getPuppeteer.js      | 45 ++++++++++++++++++++++-
 test/src/browser/commands/getPuppeteer.js | 24 ++++++------
 2 files changed, 55 insertions(+), 14 deletions(-)

diff --git a/src/browser/commands/getPuppeteer.js b/src/browser/commands/getPuppeteer.js
index 65c40c3c4..abe20c954 100644
--- a/src/browser/commands/getPuppeteer.js
+++ b/src/browser/commands/getPuppeteer.js
@@ -2,10 +2,51 @@
 
 const urljoin = require("url-join");
 
+const PUPPETEER_REJECT_TIMEOUT = 5000;
+
 module.exports.default = browser => {
     const { publicAPI: session, config } = browser;
 
-    if (!config.browserWSEndpoint || !session.getPuppeteer) {
+    if (!session.getPuppeteer) {
+        return;
+    }
+
+    session.overwriteCommand("getPuppeteer", async origGetPuppeteer => {
+        return new Promise((resolve, reject) => {
+            let isSettled = false;
+            let rejectTimeout;
+
+            origGetPuppeteer()
+                .then(puppeteer => {
+                    if (!isSettled) {
+                        resolve(puppeteer);
+                    }
+                })
+                .catch(error => {
+                    if (!isSettled) {
+                        reject(error);
+                    }
+                })
+                .finally(() => {
+                    isSettled = true;
+                    clearTimeout(rejectTimeout);
+                });
+
+            rejectTimeout = setTimeout(() => {
+                if (isSettled) {
+                    return;
+                }
+
+                isSettled = true;
+
+                browser.markAsBroken();
+
+                reject(new Error(`Unable to establish a CDP connection in ${PUPPETEER_REJECT_TIMEOUT} ms`));
+            }, PUPPETEER_REJECT_TIMEOUT);
+        });
+    });
+
+    if (!config.browserWSEndpoint) {
         return;
     }
 
@@ -17,7 +58,7 @@ module.exports.default = browser => {
         setCdpEndpoint(session.capabilities, newBrowserWSEndpoint);
 
         try {
-            return origGetPuppeteer();
+            return await origGetPuppeteer();
         } finally {
             setCdpEndpoint(session.capabilities, prevBrowserWSEndpoint);
         }
diff --git a/test/src/browser/commands/getPuppeteer.js b/test/src/browser/commands/getPuppeteer.js
index bda6c25c3..2d2e20d7d 100644
--- a/test/src/browser/commands/getPuppeteer.js
+++ b/test/src/browser/commands/getPuppeteer.js
@@ -21,23 +21,21 @@ describe('"getPuppeteer" command', () => {
         return browser.init({ sessionId: session.sessionId, sessionCaps: session.capabilities });
     };
 
-    describe("should not overwrite command if", () => {
-        it('"browserWSEndpoint" is not specified', async () => {
-            const session = mkSessionStub_();
+    it('should not overwrite command if command "getPuppeteer" does not exist in browser', async () => {
+        const session = mkSessionStub_();
+        session.getPuppeteer = undefined;
 
-            await initBrowser_({ session });
+        await initBrowser_();
 
-            assert.neverCalledWith(session.overwriteCommand, "getPuppeteer");
-        });
+        assert.neverCalledWith(session.overwriteCommand, "getPuppeteer");
+    });
 
-        it('command "getPuppeteer" does not exist in browser', async () => {
-            const session = mkSessionStub_();
-            session.getPuppeteer = undefined;
+    it('should only overwrite timeouts wrapper if "browserWSEndpoint" is not specified', async () => {
+        const session = mkSessionStub_();
 
-            await initBrowser_();
+        await initBrowser_({ session });
 
-            assert.neverCalledWith(session.overwriteCommand, "getPuppeteer");
-        });
+        assert.equal(session.overwriteCommand.withArgs("getPuppeteer").callCount, 1);
     });
 
     it("should overwrite command", async () => {
@@ -74,6 +72,8 @@ describe('"getPuppeteer" command', () => {
 
                 origGetPuppeteerFn.callsFake(() => {
                     capsBeforeReset = _.cloneDeep(session.capabilities);
+
+                    return Promise.resolve();
                 });
 
                 const browser = mkBrowser_({ browserWSEndpoint: "ws://new.endpoint/devtools" });