Fixed #1043 Fixed #1048 - correct and consistent page callback wrapping

lib/page-object/command-wrapper.js

@@ -72,7 +72,7 @@ module.exports = new (function() {
       var prevLocateStrategy = parent.client.locateStrategy;
 
       if (isElementCommand) {
-        var firstArg, desiredStrategy, callback;
+        var firstArg, desiredStrategy;
         var elementOrSectionName = args.shift();
         var getter = (isChaiAssertion && commandName === 'section') ? getSection : getElement;
         var elementOrSection = getter(parent, elementOrSectionName);
@@ -89,14 +89,25 @@ module.exports = new (function() {
         setLocateStrategy(parent.client, desiredStrategy);
         args.unshift(firstArg);
 
-        if (typeof args[args.length-1] === 'function') {
-          callback = args.pop();
-          args.push(function() {
+        // if a callback is being used with this command, wrap it in
+        // a function that allows us to restore the locate strategy
+        // to its original value before the callback is called
+
+        var callbackIndex = findCallbackIndex(args);
+        if (callbackIndex !== -1) {
+          var originalCallback = args[callbackIndex];
+
+          args[callbackIndex] = function callbackWrapper() {
+
+            // restore the locate strategy directly through client.locateStrategy.
+            // setLocateStrategy() can't be used since it uses the api commands
+            // which get added to the command queue and will not update the
+            // strategy in time for the callback which is getting immediately
+            // called after
+
             parent.client.locateStrategy = prevLocateStrategy;
-            if (callback) {
-              callback.apply(parent.client, arguments);
-            }
-          });
+            originalCallback.apply(parent.client.api, arguments);
+          };
         }
       }
 
@@ -108,6 +119,38 @@ module.exports = new (function() {
     };
   }
 
+  /**
+   * Identifies the location of a callback function within an arguments array.
+   *
+   * @param {Array} args Arguments array in which to find the location of a callback.
+   * @returns {number} Index location of the callback in the args array. If not found, -1 is returned.
+   */
+  function findCallbackIndex(args) {
+
+    if (!args) {
+      return -1;
+    }
+
+    // callbacks will usually be the last argument. waitfor methods allow an additional
+    // message argument to follow the callback which will also need to be checked for.
+
+    // last argument
+
+    var index = args.length - 1;
+    if (typeof args[index] === 'function') {
+      return index;
+    }
+
+    // second to last argument (waitfor calls)
+
+    index--;
+    if (typeof args[index] === 'function') {
+      return index;
+    }
+
+    return -1;
+  }
+
   /**
    * Retrieves an array of ancestors of the supplied element. The last element in the array is the element object itself
    *

test/src/page-object/testPageObjectCommands.js

@@ -49,6 +49,43 @@ module.exports = {
       this.client.start();
     },
 
+    testPageObjectCallbackContext : function(done) {
+      var api = this.client.api;
+      var page = api.page.simplePageObj();
+
+      page
+        .waitForElementPresent('#weblogin', 1000, true, function callback(result) {
+          assert.equal(this, api, 'page callback context using selector should equal api');
+        })
+        .waitForElementPresent('#weblogin', 1000, true, function callback(result) {
+          assert.equal(this, api, 'page callback context using selector with message should equal api');
+        }, 'Test sample message')
+        .waitForElementPresent('@loginCss', 1000, true, function callback(result) {
+          assert.equal(this, api, 'page callback context using element should equal api');
+        })
+        .waitForElementPresent('@loginCss', 1000, true, function callback(result) {
+          assert.equal(this, api, 'page callback context using element with message should equal api');
+          done();
+        }, 'Test sample message');
+
+      this.client.start();
+    },
+
+    testPageObjectLocateStrategy : function(done) {
+      var client = this.client;
+      var page = client.api.page.simplePageObj();
+
+      assert.equal(client.locateStrategy, 'css selector', 'locateStrategy should default to css selector');
+
+      page
+        .waitForElementPresent('@loginXpath', 1000, true, function callback(result) {
+          assert.equal(client.locateStrategy, 'css selector', 'locateStrategy should restore to previous css selector in callback when using xpath element');
+          done();
+        });
+
+      this.client.start();
+    },
+
     testPageObjectElementRecursion : function(done) {
       MockServer.addMock({
         'url' : '/wd/hub/session/1352110219202/element/1/click',