From 6c8920adddb373b463259c3e6c14cb3c49ecbf2b Mon Sep 17 00:00:00 2001
From: Kazuaki Matsuo <fly.49.89.over@gmail.com>
Date: Thu, 28 Mar 2024 01:34:14 -0700
Subject: [PATCH] chore: wait for wda start in sim as well for preinstalled wda
 start (#876)

* chore: wait for wda start in sim as well for preinstalled wda start

* move to inner func

* modify a bit

* remove redundant new line

* add wait for logic in getStatus

* remoe a new line

* use waitForCondition in getSttus

* modify error handling

* tweak error log

* tweak log further

* Update webdriveragent.js
---
 lib/webdriveragent.js | 73 +++++++++++++++++++++++++++++--------------
 1 file changed, 50 insertions(+), 23 deletions(-)

diff --git a/lib/webdriveragent.js b/lib/webdriveragent.js
index 885291e76..e374eb6ad 100644
--- a/lib/webdriveragent.js
+++ b/lib/webdriveragent.js
@@ -190,21 +190,53 @@ class WebDriverAgent {
    *   }
    * }
    *
-   * @return {Promise<any?>} State Object
+   * @param {number} [timeoutMs=0] If the given timeoutMs is zero or negative number,
+   *                               this function will return the response of `/status` immediately. If the given timeoutMs,
+   *                               this function will try to get the response of `/status` up to the timeoutMs.
+   * @return {Promise<import('@appium/types').StringRecord|null>} State Object
+   * @throws {Error} If there was an error within timeoutMs timeout.
+   *                 No error is raised if zero or negative number for the timeoutMs.
    */
-  async getStatus () {
+  async getStatus (timeoutMs = 0) {
     const noSessionProxy = new NoSessionProxy({
       server: this.url.hostname,
       port: this.url.port,
       base: this.basePath,
       timeout: 3000,
     });
+
+    const sendGetStatus = async () => await /** @type import('@appium/types').StringRecord */ (noSessionProxy.command('/status', 'GET'));
+
+    if (_.isNil(timeoutMs) || timeoutMs <= 0) {
+      try {
+        return await sendGetStatus();
+      } catch (err) {
+        this.log.debug(`WDA is not listening at '${this.url.href}'. Original error:: ${err.message}`);
+        return null;
+      }
+    }
+
+    let lastError = null;
+    let status = null;
     try {
-      return await noSessionProxy.command('/status', 'GET');
+      await waitForCondition(async () => {
+        try {
+          status = await sendGetStatus();
+          return true;
+        } catch (err) {
+          lastError = err;
+        }
+        return false;
+      }, {
+        waitMs: timeoutMs,
+        intervalMs: 300,
+      });
     } catch (err) {
-      this.log.debug(`WDA is not listening at '${this.url.href}'`);
-      return null;
+      this.log.debug(`Failed to get the status endpoint in ${timeoutMs} ms. ` +
+        `The last error while accessing ${this.url.href}: ${lastError}. Original error:: ${err.message}.`);
+      throw new Error(`WDA was not ready in ${timeoutMs} ms.`);
     }
+    return status;
   }
 
   /**
@@ -317,28 +349,14 @@ class WebDriverAgent {
     await this.device.devicectl.launchApp(
       this.bundleIdForXctest, { env, terminateExisting: true }
     );
-
-    // Launching app via decictl does not wait for the app start.
-    // We should wait for the app start by ourselves.
-    try {
-      await waitForCondition(async () => !_.isNull(await this.getStatus()), {
-        waitMs: this.wdaLaunchTimeout,
-        intervalMs: 300,
-      });
-
-    } catch (err) {
-      throw new Error(
-        `Failed to start the preinstalled WebDriverAgent in ${this.wdaLaunchTimeout} ms. ` +
-        `The WebDriverAgent might not be properly built or the device might be locked. ` +
-        `The 'appium:wdaLaunchTimeout' capability modifies the timeout.`
-      );
-    }
   }
 
   /**
    * Launch WDA with preinstalled package without xcodebuild.
    * @param {string} sessionId Launch WDA and establish the session with this sessionId
-   * @return {Promise<any?>} State Object
+   * @return {Promise<import('@appium/types').StringRecord|null>} State Object
+   * @throws {Error} If there was an error within timeoutMs timeout.
+   *                 No error is raised if zero or negative number for the timeoutMs.
    */
   async launchWithPreinstalledWDA(sessionId) {
     const xctestEnv = {
@@ -371,7 +389,16 @@ class WebDriverAgent {
     }
 
     this.setupProxies(sessionId);
-    const status = await this.getStatus();
+    let status;
+    try {
+      status = await this.getStatus(this.wdaLaunchTimeout);
+    } catch (err) {
+      throw new Error(
+        `Failed to start the preinstalled WebDriverAgent in ${this.wdaLaunchTimeout} ms. ` +
+        `The WebDriverAgent might not be properly built or the device might be locked. ` +
+        `The 'appium:wdaLaunchTimeout' capability modifies the timeout.`
+      );
+    }
     this.started = true;
     return status;
   }