Added support for defining a custom return function for custom command

Author: beatfactor

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;

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 {

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();
+  });
+});

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++;
+    });
+  }
+};

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
+    }));
+  });
+
 });