From bb6ef6e7510b70504778113b47210b0a0e045458 Mon Sep 17 00:00:00 2001
From: saikrishna321 <saikrishna321@yahoo.com>
Date: Mon, 20 Jan 2025 16:11:53 +0530
Subject: [PATCH] Fix review comments

Co-authored-by: SrinivasanTarget <srinivasan.sekar1990@gmail.com>
---
 lib/driver.js | 54 ++++++++++++++++++++++++++++++---------------------
 1 file changed, 32 insertions(+), 22 deletions(-)

diff --git a/lib/driver.js b/lib/driver.js
index 8fbd8e968..b00e989b6 100644
--- a/lib/driver.js
+++ b/lib/driver.js
@@ -467,29 +467,12 @@ export class XCUITestDriver extends BaseDriver {
       // ensure WDA gets our defaults instead of whatever its own might be
       await this.updateSettings(wdaSettings);
 
-      if(this.opts.mjpegServerPort) {
-        this.log.debug(
-          `MJPEG broadcasting requested, forwarding MJPEG server port ${this.opts.mjpegServerPort} ` +
-            `to local port ${this.opts.mjpegServerPort}`
-        );
-        try {
-          await DEVICE_CONNECTIONS_FACTORY.requestConnection(caps.udid, this.opts.mjpegServerPort, {
-            devicePort: this.opts.mjpegServerPort,
-            usePortForwarding: true,
-          });
-        } catch {
-          this.log.warn(
-            `Cannot forward the local port ${this.opts.mjpegServerPort} to ${this.opts.mjpegServerPort} ` +
-              `on the device ${caps.udid}.`,
-          );
-        }
-      }
-      // turn on mjpeg stream reading if requested
-      if (this.opts.mjpegScreenshotUrl) {
-        this.log.info(`Starting MJPEG stream reading URL: '${this.opts.mjpegScreenshotUrl}'`);
-        this.mjpegStream = new mjpeg.MJpegStream(this.opts.mjpegScreenshotUrl);
-        await this.mjpegStream.start();
+      if (this.isRealDevice() && this.opts.mjpegServerPort) {
+        await this.allocateMjpegServerPort(caps);
       }
+
+      await this.startMjpegStream();
+
       return /** @type {[string, import('@appium/types').DriverCaps<XCUITestDriverConstraints>]} */ ([
         sessionId,
         caps,
@@ -501,6 +484,33 @@ export class XCUITestDriver extends BaseDriver {
     }
   }
 
+  async startMjpegStream() {
+      // turn on mjpeg stream reading if requested
+      if (this.opts.mjpegScreenshotUrl) {
+        this.log.info(`Starting MJPEG stream reading URL: '${this.opts.mjpegScreenshotUrl}'`);
+        this.mjpegStream = new mjpeg.MJpegStream(this.opts.mjpegScreenshotUrl);
+        await this.mjpegStream.start();
+      }
+  }
+
+  async allocateMjpegServerPort(caps) {
+    this.log.debug(
+      `MJPEG broadcasting requested, forwarding MJPEG server port ${this.opts.mjpegServerPort} ` +
+      `to local port ${this.opts.mjpegServerPort}`
+    );
+    try {
+      await DEVICE_CONNECTIONS_FACTORY.requestConnection(caps.udid, this.opts.mjpegServerPort, {
+        devicePort: this.opts.mjpegServerPort,
+        usePortForwarding: true,
+      });
+    } catch {
+       throw new Error(
+        `Cannot forward the local port ${this.opts.mjpegServerPort} to ${this.opts.mjpegServerPort} ` +
+        `on the device ${caps.udid}.`
+      );
+    }
+  }
+
   /**
    * Returns the default URL for Safari browser
    * @returns {string} The default URL