Support functions for message in driver.wait()

javascript/node/selenium-webdriver/lib/webdriver.js

@@ -222,6 +222,15 @@ function fromWireValue(driver, value) {
   return value;
 }
 
+/**
+ * Resolves a wait message from either a function or a string.
+ * @param {(string|Function)=} message An optional message to use if the wait times out.
+ * @return {string} The resolved message
+ */
+function resolveWaitMessage(message) {
+  return message ? `${typeof message === 'function' ? message() : message}\n` : '';
+}
+
 
 /**
  * Structural interface for a WebDriver client.
@@ -435,7 +444,7 @@ class IWebDriver {
    *     evaluate as a condition.
    * @param {number=} timeout The duration in milliseconds, how long to wait
    *     for the condition to be true.
-   * @param {string=} message An optional message to use if the wait times out.
+   * @param {(string|Function)=} message An optional message to use if the wait times out.
    * @param {number=} pollTimeout The duration in milliseconds, how long to
    *     wait between polling the condition.
    * @return {!(IThenable<T>|WebElementPromise)} A promise that will be
@@ -787,11 +796,13 @@ class WebDriver {
         let start = Date.now();
         let timer = setTimeout(function() {
           timer = null;
-          reject(
-              new error.TimeoutError(
-                  (message ? `${message}\n` : '')
-                      + 'Timed out waiting for promise to resolve after '
-                      + (Date.now() - start) + 'ms'));
+          try {
+            let timeoutMessage = resolveWaitMessage(message);
+            reject(new error.TimeoutError(`${timeoutMessage}Timed out waiting for promise to resolve after ${Date.now() - start}ms`));
+          }
+          catch (ex) {
+            reject(new error.TimeoutError(`${ex.message}\nTimed out waiting for promise to resolve after ${Date.now() - start}ms`));
+          }
         }, timeout);
         const clearTimer = () => timer && clearTimeout(timer);
 
@@ -838,10 +849,13 @@ class WebDriver {
           if (!!value) {
             resolve(value);
           } else if (timeout && elapsed >= timeout) {
-            reject(
-                new error.TimeoutError(
-                  (message ? `${message}\n` : '')
-                        + `Wait timed out after ${elapsed}ms`));
+            try {
+              let timeoutMessage = resolveWaitMessage(message);
+              reject(new error.TimeoutError(`${timeoutMessage}Wait timed out after ${elapsed}ms`));
+            }
+            catch (ex) {
+              reject(new error.TimeoutError(`${ex.message}\nWait timed out after ${elapsed}ms`));
+            }
           } else {
             setTimeout(pollCondition, pollTimeout);
           }

javascript/node/selenium-webdriver/test/lib/webdriver_test.js

@@ -1040,6 +1040,48 @@ describe('WebDriver', function() {
           .then(fail, assertIsStubError);
     });
 
+    it('supports message function if condition exceeds timeout', function() {
+      let executor = new FakeExecutor();
+      let driver = executor.createDriver();
+      let message = () => 'goes boom';
+      return driver.wait(() => false, 0.001, message)
+        .then(fail, (e) => {
+          assert.ok(/^goes boom\nWait timed out after \d+ms$/.test(e.message));
+        })
+    });
+
+    it('handles if the message function throws an error after a condition exceeds timeout', function() {
+      let executor = new FakeExecutor();
+      let driver = executor.createDriver();
+      let message = () => { throw new Error('message function error') };
+      return driver.wait(() => false, 0.001, message)
+        .then(fail, (e) => {
+          assert.ok(/^message function error\nWait timed out after \d+ms$/.test(e.message));
+        })
+    });
+
+    it('supports message function if condition returns a rejected promise', function() {
+      let executor = new FakeExecutor();
+      let driver = executor.createDriver();
+      let condition = new Promise((res) => setTimeout(res, 100));
+      let message = () => 'goes boom';
+      return driver.wait(condition, 1, message)
+        .then(fail, (e) => {
+          assert.ok(/^goes boom\nTimed out waiting for promise to resolve after \d+ms$/.test(e.message));
+        });
+    });
+
+    it('handles if the message function returns an error after a rejected promise', function() {
+      let executor = new FakeExecutor();
+      let driver = executor.createDriver();
+      let condition = new Promise((res) => setTimeout(res, 100));
+      let message = () => { throw new Error('message function error') };
+      return driver.wait(condition, 1, message)
+        .then(fail, (e) => {
+          assert.ok(/^message function error\nTimed out waiting for promise to resolve after \d+ms$/.test(e.message));
+        });
+    });
+
     it('waits forever on a zero timeout', function() {
       let done = false;
       setTimeout(() => done = true, 150);