Fix waitForFunction() (#294)

* Fix waitForFunction()

This is essentially a re-implementation of `waitForFunction()` since the
previous version was broken in many ways. This required changes in
`injected_script.js` where I hopefully didn't break anything else.
I tried to split it into several commits but the changes were too messy
and it would be very difficult at this point.

In summary this:
- Fixes `waitForFunction()` option parsing.
- Ensures all three polling variants work: `requestAnimationFrame`, DOM
  mutation and interval.
- Ensures we properly time out in all cases and return an informative
  error message.
- Makes `waitForFunction()` an async method.

Fixes #291

* Resolve promise on truthy value, return JSHandle

Resolves #294 (comment)

* Add more tests

Resolves #294 (comment)

* Update waitForFunction example

* Add comments to waitForFunction eval calls

Author: Ivan Mirić

api/frame.go

@@ -67,7 +67,7 @@ type Frame interface {
 	Type(selector string, text string, opts goja.Value)
 	Uncheck(selector string, opts goja.Value)
 	URL() string
-	WaitForFunction(pageFunc goja.Value, opts goja.Value, args ...goja.Value) JSHandle
+	WaitForFunction(pageFunc, opts goja.Value, args ...goja.Value) *goja.Promise
 	WaitForLoadState(state string, opts goja.Value)
 	WaitForNavigation(opts goja.Value) Response
 	WaitForSelector(selector string, opts goja.Value) ElementHandle

api/js_handle.go

@@ -20,7 +20,10 @@
 
 package api
 
-import "github.com/dop251/goja"
+import (
+	cdpruntime "github.com/chromedp/cdproto/runtime"
+	"github.com/dop251/goja"
+)
 
 // JSHandle is the interface of an in-page JS object.
 type JSHandle interface {
@@ -31,4 +34,5 @@ type JSHandle interface {
 	GetProperties() map[string]JSHandle
 	GetProperty(propertyName string) JSHandle
 	JSONValue() goja.Value
+	ObjectID() cdpruntime.RemoteObjectID
 }

api/page.go

@@ -88,7 +88,7 @@ type Page interface {
 	Video() Video
 	ViewportSize() map[string]float64
 	WaitForEvent(event string, optsOrPredicate goja.Value) interface{}
-	WaitForFunction(pageFunc goja.Value, arg goja.Value, opts goja.Value) JSHandle
+	WaitForFunction(fn, opts goja.Value, args ...goja.Value) *goja.Promise
 	WaitForLoadState(state string, opts goja.Value)
 	WaitForNavigation(opts goja.Value) Response
 	WaitForRequest(urlOrPredicate, opts goja.Value) Request

common/element_handle.go

@@ -672,7 +672,7 @@ func (h *ElementHandle) waitForElementState(apiCtx context.Context, states []str
 	}
 	result, err := h.evalWithScript(apiCtx, opts, fn, states, timeout.Milliseconds())
 	if err != nil {
-		return false, err
+		return false, errorFromDOMError(err.Error())
 	}
 
 	value := result.(goja.Value)
@@ -1249,7 +1249,7 @@ func (h *ElementHandle) Uncheck(opts goja.Value) {
 }
 
 func (h *ElementHandle) WaitForElementState(state string, opts goja.Value) {
-	parsedOpts := NewElementHandleWaitForElementStateOptions(time.Duration(h.frame.manager.timeoutSettings.timeout()) * time.Second)
+	parsedOpts := NewElementHandleWaitForElementStateOptions(h.defaultTimeout())
 	err := parsedOpts.Parse(h.ctx, opts)
 	if err != nil {
 		k6Throw(h.ctx, "cannot parse element wait for state options: %w", err)
@@ -1458,7 +1458,7 @@ func retryPointerAction(
 
 func errorFromDOMError(derr string) error {
 	// return the same sentinel error value for the timed out err
-	if derr == "error:timeout" {
+	if strings.Contains(derr, "timed out") {
 		return ErrTimedOut
 	}
 	if s := "error:expectednode:"; strings.HasPrefix(derr, s) {

common/element_handle_test.go

@@ -18,7 +18,7 @@ func TestErrorFromDOMError(t *testing.T) {
 		sentinel bool // if it returns the same error value
 		want     error
 	}{
-		{in: "error:timeout", want: ErrTimedOut, sentinel: true},
+		{in: "timed out", want: ErrTimedOut, sentinel: true},
 		{in: "error:notconnected", want: errors.New("element is not attached to the DOM")},
 		{in: "error:expectednode:anything", want: errors.New("expected node but got anything")},
 		{in: "nonexistent error", want: errors.New("nonexistent error")},

common/frame.go

@@ -567,13 +567,12 @@ func (f *Frame) waitForExecutionContext(world executionWorld) {
 }
 
 func (f *Frame) waitForFunction(
-	apiCtx context.Context, world executionWorld, predicateFn string,
-	polling PollingType, interval int64, timeout time.Duration,
-	args ...interface{},
-) (interface{}, error) {
+	apiCtx context.Context, world executionWorld, js string,
+	polling interface{}, timeout time.Duration, args ...interface{},
+) (*goja.Promise, error) {
 	f.log.Debugf(
 		"Frame:waitForFunction",
-		"fid:%s furl:%q world:%s pt:%s timeout:%s",
+		"fid:%s furl:%q world:%s poll:%s timeout:%s",
 		f.ID(), f.URL(), world, polling, timeout)
 
 	f.waitForExecutionContext(world)
@@ -592,23 +591,53 @@ func (f *Frame) waitForFunction(
 
 	pageFn := `
 		(injected, predicate, polling, timeout, ...args) => {
-			return injected.waitForPredicateFunction(predicate, polling, timeout, ...args);
+			const fn = (...args) => {
+				return predicate(...args) || injected.continuePolling;
+			}
+			return injected.waitForPredicateFunction(fn, polling, timeout, ...args);
 		}
 	`
-	opts := evalOptions{
-		forceCallable: true,
-		returnByValue: false,
-	}
-	result, err := execCtx.eval(
-		apiCtx, opts, pageFn, append([]interface{}{
-			injected,
-			predicateFn,
-			polling,
-		}, args...)...)
-	if err != nil {
-		return nil, fmt.Errorf("frame cannot wait for function: %w", err)
-	}
-	return result, nil
+
+	cb := f.vu.RegisterCallback()
+	rt := f.vu.Runtime()
+	promise, resolve, reject := rt.NewPromise()
+
+	go func() {
+		// First evaluate the predicate function itself to get its handle.
+		opts := evalOptions{forceCallable: false, returnByValue: false}
+		handle, err := execCtx.eval(apiCtx, opts, js)
+		if err != nil {
+			cb(func() error {
+				reject(fmt.Errorf("waitForFunction promise rejected: %w", err))
+				return nil
+			})
+			return
+		}
+
+		// Then evaluate the injected function call, passing it the predicate
+		// function handle and the rest of the arguments.
+		opts = evalOptions{forceCallable: true, returnByValue: false}
+		result, err := execCtx.eval(
+			apiCtx, opts, pageFn, append([]interface{}{
+				injected,
+				handle,
+				polling,
+				timeout.Milliseconds(), // The JS value is in ms integers
+			}, args...)...)
+		if err != nil {
+			cb(func() error {
+				reject(fmt.Errorf("waitForFunction promise rejected: %w", err))
+				return nil
+			})
+			return
+		}
+		cb(func() error {
+			resolve(result)
+			return nil
+		})
+	}()
+
+	return promise, nil
 }
 
 func (f *Frame) waitForSelectorRetry(
@@ -1471,13 +1500,13 @@ func (f *Frame) setURL(url string) {
 }
 
 // WaitForFunction waits for the given predicate to return a truthy value.
-func (f *Frame) WaitForFunction(fn goja.Value, opts goja.Value, jsArgs ...goja.Value) api.JSHandle {
+func (f *Frame) WaitForFunction(fn goja.Value, opts goja.Value, jsArgs ...goja.Value) *goja.Promise {
 	f.log.Debugf("Frame:WaitForFunction", "fid:%s furl:%q", f.ID(), f.URL())
 
 	parsedOpts := NewFrameWaitForFunctionOptions(f.defaultTimeout())
 	err := parsedOpts.Parse(f.ctx, opts)
 	if err != nil {
-		k6Throw(f.ctx, "failed parsing options: %w", err)
+		k6Throw(f.ctx, "error parsing waitForFunction options: %w", err)
 	}
 
 	f.executionContextMu.RLock()
@@ -1491,15 +1520,23 @@ func (f *Frame) WaitForFunction(fn goja.Value, opts goja.Value, jsArgs ...goja.V
 
 	args := make([]interface{}, 0, len(jsArgs))
 	for _, a := range jsArgs {
-		args = append(args, a.Export())
+		if a != nil {
+			args = append(args, a.Export())
+		}
 	}
 
-	handle, err := f.waitForFunction(f.ctx, utilityWorld, js,
-		parsedOpts.Polling, parsedOpts.Interval, parsedOpts.Timeout, args...)
+	var polling interface{} = parsedOpts.Polling
+	if parsedOpts.Polling == PollingInterval {
+		polling = parsedOpts.Interval
+	}
+
+	promise, err := f.waitForFunction(f.ctx, mainWorld, js,
+		polling, parsedOpts.Timeout, args...)
 	if err != nil {
 		k6Throw(f.ctx, "%w", err)
 	}
-	return handle.(api.JSHandle)
+
+	return promise
 }
 
 // WaitForLoadState waits for the given load state to be reached.

common/frame_options.go

@@ -605,22 +605,31 @@ func NewFrameWaitForFunctionOptions(defaultTimeout time.Duration) *FrameWaitForF
 	}
 }
 
+// Parse JavaScript waitForFunction options.
 func (o *FrameWaitForFunctionOptions) Parse(ctx context.Context, opts goja.Value) error {
 	rt := GetVU(ctx).Runtime()
 
 	if opts != nil && !goja.IsUndefined(opts) && !goja.IsNull(opts) {
 		opts := opts.ToObject(rt)
 		for _, k := range opts.Keys() {
+			v := opts.Get(k)
 			switch k {
 			case "timeout":
-				o.Timeout = time.Duration(opts.Get(k).ToInteger()) * time.Millisecond
+				o.Timeout = time.Duration(v.ToInteger()) * time.Millisecond
 			case "polling":
-				switch opts.Get(k).ExportType().Kind() {
+				switch v.ExportType().Kind() { //nolint: exhaustive
 				case reflect.Int64:
 					o.Polling = PollingInterval
-					o.Interval = opts.Get(k).ToInteger()
+					o.Interval = v.ToInteger()
+				case reflect.String:
+					if p, ok := pollingTypeToID[v.ToString().String()]; ok {
+						o.Polling = p
+						break
+					}
+					fallthrough
 				default:
-					o.Polling = PollingRaf
+					return fmt.Errorf("wrong polling option value: %q; "+
+						`possible values: "raf", "mutation" or number`, v)
 				}
 			}
 		}

common/js/injected_script.js

@@ -49,17 +49,17 @@ const eventType = new Map([
     ['mouseup', 'mouse'],
     ['mouseleave', 'mouse'],
     ['mousewheel', 'mouse'],
-  
+
     ['keydown', 'keyboard'],
     ['keyup', 'keyboard'],
     ['keypress', 'keyboard'],
     ['textInput', 'keyboard'],
-  
+
     ['touchstart', 'touch'],
     ['touchmove', 'touch'],
     ['touchend', 'touch'],
     ['touchcancel', 'touch'],
-  
+
     ['pointerover', 'pointer'],
     ['pointerout', 'pointer'],
     ['pointerenter', 'pointer'],
@@ -70,10 +70,10 @@ const eventType = new Map([
     ['pointercancel', 'pointer'],
     ['gotpointercapture', 'pointer'],
     ['lostpointercapture', 'pointer'],
-  
+
     ['focus', 'focus'],
     ['blur', 'focus'],
-  
+
     ['drag', 'drag'],
     ['dragstart', 'drag'],
     ['dragend', 'drag'],
@@ -138,6 +138,7 @@ class InjectedScript {
     constructor() {
         this._replaceRafWithTimeout = false;
         this._stableRafCount = 10;
+        this.continuePolling = continuePolling;
         this._queryEngines = {
             'css': new CSSQueryEngine(),
             'text': new TextQueryEngine(),
@@ -587,30 +588,42 @@ class InjectedScript {
     }
 
     async waitForPredicateFunction(predicate, polling, timeout, ...args) {
-        predicate();
         let timedOut = false;
-        if (timeout !== undefined || timeout !== null ) setTimeout(() => (timedOut = true), timeout);
+        let timeoutPoll = null;
+        if (timeout !== undefined || timeout !== null) {
+          setTimeout(() => {
+            timedOut = true;
+            if (timeoutPoll) timeoutPoll();
+          }, timeout);
+        }
         if (polling === 'raf') return await pollRaf();
         if (polling === 'mutation') return await pollMutation();
         if (typeof polling === 'number') return await pollInterval(polling);
 
         async function pollMutation() {
-            const success = await predicate(...args);
-            if (success) return Promise.resolve(success);
+            const success = predicate(...args);
+            if (success !== continuePolling) return Promise.resolve(success);
 
-            let fulfill;
-            const result = new Promise((x) => (fulfill = x));
+            let resolve, reject;
+            const result = new Promise((res, rej) => {
+              resolve = res;
+              reject = rej;
+            });
             const observer = new MutationObserver(async () => {
                 if (timedOut) {
                     observer.disconnect();
-                    fulfill("error:timeout");
+                    reject(`timed out after ${timeout}ms`);
                 }
                 const success = predicate(...args);
                 if (success !== continuePolling) {
                     observer.disconnect();
-                    fulfill(success);
+                    resolve(success);
                 }
             });
+            timeoutPoll = () => {
+              observer.disconnect();
+              reject(`timed out after ${timeout}ms`);
+            };
             observer.observe(document, {
                 childList: true,
                 subtree: true,
@@ -620,35 +633,41 @@ class InjectedScript {
         }
 
         async function pollRaf() {
-            let fulfill;
-            const result = new Promise((x) => (fulfill = x));
+            let resolve, reject;
+            const result = new Promise((res, rej) => {
+              resolve = res;
+              reject = rej;
+            });
             await onRaf();
             return result;
 
             async function onRaf() {
                 if (timedOut) {
-                    fulfill("error:timeout");
+                    reject(`timed out after ${timeout}ms`);
                     return;
                 }
                 const success = predicate(...args);
-                if (success !== continuePolling) fulfill(success);
+                if (success !== continuePolling) resolve(success);
                 else requestAnimationFrame(onRaf);
             }
         }
 
         async function pollInterval(pollInterval) {
-            let fulfill;
-            const result = new Promise((x) => (fulfill = x));
+            let resolve, reject;
+            const result = new Promise((res, rej) => {
+              resolve = res;
+              reject = rej;
+            });
             await onTimeout();
             return result;
 
             async function onTimeout() {
                 if (timedOut) {
-                    fulfill("error:timeout");
+                    reject(`timed out after ${timeout}ms`);
                     return;
                 }
                 const success = predicate(...args);
-                if (success !== continuePolling) fulfill(success);
+                if (success !== continuePolling) resolve(success);
                 else setTimeout(onTimeout, pollInterval);
             }
         }
@@ -760,4 +779,4 @@ class InjectedScript {
 
         return this.waitForPredicateFunction(predicate, polling, timeout, ...args);
     }
-}
+}

common/js_handle.go

@@ -168,3 +168,8 @@ func (h *BaseJSHandle) JSONValue() goja.Value {
 	}
 	return res
 }
+
+// ObjectID returns the remote object ID.
+func (h *BaseJSHandle) ObjectID() runtime.RemoteObjectID {
+	return h.remoteObject.ObjectID
+}