From fe9a010a4787aa6395065eea3dc39e92116bd314 Mon Sep 17 00:00:00 2001
From: Andrei Rusu <andrei.rusu@beatfactor.net>
Date: Fri, 29 Dec 2023 13:45:24 +0100
Subject: [PATCH] Added support for defining a custom return function for
 custom command

---
 lib/api/_loaders/_base-loader.js              |  8 ++++-
 lib/api/web-element/scoped-element.js         |  4 +++
 .../testUsingCommandReturnFn.js               | 14 +++++++++
 .../returnFn/customCommandReturnFn.js         | 17 +++++++++++
 .../testCustomCommandAutoInvoke.js            | 30 ++++++++++++++++++-
 5 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 test/apidemos/custom-commands/testUsingCommandReturnFn.js
 create mode 100644 test/extra/commands/returnFn/customCommandReturnFn.js

diff --git a/lib/api/_loaders/_base-loader.js b/lib/api/_loaders/_base-loader.js
index 14c2a4c7d1..d36abbef91 100644
--- a/lib/api/_loaders/_base-loader.js
+++ b/lib/api/_loaders/_base-loader.js
@@ -467,7 +467,13 @@ class BaseLoader extends EventEmitter {
           BaseLoader.lastDeferred.reject(err);
         });
 
-        return node.deferred.promise;
+        if (!this.module?.returnFn) {
+          return node.deferred.promise;
+        }
+      }
+
+      if (this.module?.returnFn) {
+        return this.module.returnFn(node, apiToReturn || api);
       }
 
       return apiToReturn || api;
diff --git a/lib/api/web-element/scoped-element.js b/lib/api/web-element/scoped-element.js
index 1514a91048..b127f11041 100644
--- a/lib/api/web-element/scoped-element.js
+++ b/lib/api/web-element/scoped-element.js
@@ -147,6 +147,10 @@ class ScopedWebElement {
 
   async findElementAction({parentElement, condition, index, timeout, retryInterval, abortOnFailure, suppressNotFoundErrors}) {
     const createAction = () => async ({args}) => {
+      if ((args[0] instanceof Promise) && !args[0]['@nightwatch_element']) {
+        args[0] = await args[0];
+      }
+
       const parentElement = args[0];
 
       try {
diff --git a/test/apidemos/custom-commands/testUsingCommandReturnFn.js b/test/apidemos/custom-commands/testUsingCommandReturnFn.js
new file mode 100644
index 0000000000..5c7f9fb31e
--- /dev/null
+++ b/test/apidemos/custom-commands/testUsingCommandReturnFn.js
@@ -0,0 +1,14 @@
+const assert = require('assert');
+
+describe('Test using custom commands with returnFn', function() {
+  before(browser => {
+    browser
+      .url('http://localhost');
+  });
+
+  it('sampleTest', browser => {
+    const result = browser.customCommandReturnFn();
+    assert.deepStrictEqual(result, {status: 0});
+    browser.end();
+  });
+});
diff --git a/test/extra/commands/returnFn/customCommandReturnFn.js b/test/extra/commands/returnFn/customCommandReturnFn.js
new file mode 100644
index 0000000000..3460cef2f3
--- /dev/null
+++ b/test/extra/commands/returnFn/customCommandReturnFn.js
@@ -0,0 +1,17 @@
+module.exports = class CustomCommandReturnFn {
+  static get returnFn() {
+    return function(node, api) {
+      api.globals.count++;
+
+      return {
+        status: 0
+      };
+    };
+  }
+
+  command() {
+    this.api.perform(function() {
+      this.globals.count++;
+    });
+  }
+};
diff --git a/test/src/apidemos/custom-commands/testCustomCommandAutoInvoke.js b/test/src/apidemos/custom-commands/testCustomCommandAutoInvoke.js
index c83c9b2471..1f5eb33194 100644
--- a/test/src/apidemos/custom-commands/testCustomCommandAutoInvoke.js
+++ b/test/src/apidemos/custom-commands/testCustomCommandAutoInvoke.js
@@ -20,7 +20,7 @@ describe('custom commands with auto invoke', function() {
     });
   });
 
-  it('custom find elements es6 async', function() {
+  it('test custom command using autoInvoke', function() {
     const testsPath = path.join(__dirname, '../../../apidemos/custom-commands/testUsingAutoInvokeCommand.js');
     Mocks.createNewW3CSession({
       testName: 'Test Using ES6 Async Custom Commands'
@@ -50,4 +50,32 @@ describe('custom commands with auto invoke', function() {
     }));
   });
 
+  it('test custom command with returnFn', function() {
+    const testsPath = path.join(__dirname, '../../../apidemos/custom-commands/testUsingCommandReturnFn.js');
+    Mocks.createNewW3CSession({
+      testName: 'Test using custom commands with returnFn'
+    });
+
+    const globals = {
+      waitForConditionPollInterval: 50,
+      waitForConditionTimeout: 120,
+      retryAssertionTimeout: 1000,
+      count: 0,
+      reporter(results) {
+        if (results.lastError) {
+          throw results.lastError;
+        }
+        assert.strictEqual(this.count, 2);
+      }
+    };
+
+    return NightwatchClient.runTests(testsPath, settings({
+      output: true,
+      silent: false,
+      selenium_host: null,
+      custom_commands_path: [path.join(__dirname, '../../../extra/commands/returnFn')],
+      globals
+    }));
+  });
+
 });
\ No newline at end of file