Fix some of the issues in the Poll action

poll.go

@@ -35,15 +35,29 @@ func (p *pollTask) Do(ctx context.Context) error {
 	if t == nil {
 		return ErrInvalidTarget
 	}
-	var frameID cdp.FrameID
-	if p.frame == nil {
-		frameID = t.cur
+	var (
+		root    *cdp.Node
+		execCtx runtime.ExecutionContextID
+		ok      bool
+	)
+	for !ok {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(5 * time.Millisecond):
+		}
+		_, root, execCtx, ok = t.ensureFrame()
+	}
+
+	fromNode := p.frame
+	if fromNode == nil {
+		fromNode = root
 	} else {
-		frameID = t.enclosingFrame(p.frame)
+		t.frameMu.RLock()
+		frameID := t.enclosingFrame(fromNode)
+		execCtx = t.execContexts[frameID]
+		t.frameMu.RUnlock()
 	}
-	t.frameMu.RLock()
-	execCtx := t.execContexts[frameID]
-	t.frameMu.RUnlock()
 
 	ea := &errAppender{args: make([]*runtime.CallArgument, 0, len(p.args)+3)}
 	ea.append(p.predicate)

poll_test.go

@@ -89,26 +89,42 @@ func TestPoll(t *testing.T) {
 			delay:  50 * time.Millisecond,
 		},
 		{
-			name:   "PollingMutation",
-			js:     "globalThis.__FOO === 1",
-			isFunc: false,
+			name: "PollingMutation",
+			js: `() => {
+				if (globalThis.__Mutation === 1){
+					return true;
+				} else {
+					globalThis.__Mutation = 1;
+					setTimeout(() => {
+						document.body.appendChild(document.createElement('div'))
+					}, 100);
+				}
+			}`,
+			isFunc: true,
 			opts:   []PollOption{WithPollingMutation(), WithPollingTimeout(200 * time.Millisecond)},
-			hash:   "#mutation",
 			delay:  50 * time.Millisecond,
 		},
 		{
 			name:   "TimeoutWithoutMutation",
-			js:     "globalThis.__FOO === 1",
+			js:     "globalThis.__Mutation === 1",
 			isFunc: false,
 			opts:   []PollOption{WithPollingMutation(), WithPollingTimeout(100 * time.Millisecond)},
 			err:    ErrPollingTimeout.Error(),
 		},
 		{
-			name:   "TimeoutBeforeMutation",
-			js:     "globalThis.__FOO === 1",
-			isFunc: false,
+			name: "TimeoutBeforeMutation",
+			js: `() => {
+				if (globalThis.__Mutation === 1){
+					return true;
+				} else {
+					globalThis.__Mutation = 1;
+					setTimeout(() => {
+						document.body.appendChild(document.createElement('div'))
+					}, 100);
+				}
+			}`,
+			isFunc: true,
 			opts:   []PollOption{WithPollingMutation(), WithPollingTimeout(50 * time.Millisecond)},
-			hash:   "#mutation",
 			err:    ErrPollingTimeout.Error(),
 		},
 		{
@@ -135,16 +151,20 @@ func TestPoll(t *testing.T) {
 				Navigate(testdataDir+"/poll.html"+test.hash),
 				action,
 			)
-			if test.err != "" {
+			if test.err == "" {
+				if err != nil {
+					t.Fatalf("got error: %v", err)
+				} else if !res {
+					t.Fatalf("got no error, but res is not true")
+				}
+
+			} else {
 				if err == nil {
 					t.Fatalf("expected err to be %q, got: %v", test.err, err)
 				} else if test.err != err.Error() {
 					t.Fatalf("want error to be %v, got: %v", test.err, err)
 				}
 			}
-			if test.err == "" && !res {
-				t.Fatalf("got no error, but res is not true")
-			}
 			if test.delay != 0 {
 				delay := time.Since(startTime)
 				if delay < test.delay {

query.go

@@ -165,30 +165,17 @@ func (s *Selector) Do(ctx context.Context) error {
 		case <-time.After(5 * time.Millisecond):
 		}
 
-		t.frameMu.RLock()
-		frame := t.frames[t.cur]
-		execCtx := t.execContexts[t.cur]
-
-		if frame == nil || execCtx == 0 {
-			// the frame hasn't loaded yet.
-			t.frameMu.RUnlock()
+		frame, root, execCtx, ok := t.ensureFrame()
+		if !ok {
 			continue
 		}
 
 		fromNode := s.fromNode
 		if fromNode == nil {
-			t.frameMu.RUnlock()
-
-			frame.RLock()
-			fromNode = frame.Root
-			frame.RUnlock()
-
-			if fromNode == nil {
-				// not root node yet?
-				continue
-			}
+			fromNode = root
 		} else {
 			frameID := t.enclosingFrame(fromNode)
+			t.frameMu.RLock()
 			execCtx = t.execContexts[frameID]
 			t.frameMu.RUnlock()
 

target.go

@@ -63,6 +63,31 @@ func (t *Target) enclosingFrame(node *cdp.Node) cdp.FrameID {
 	return node.FrameID
 }
 
+// ensureFrame ensures the top frame of this target is loaded and returns the top frame,
+// the root node and the ExecutionContextID of this top frame; otherwise, it will return
+// false as its last return value.
+func (t *Target) ensureFrame() (*cdp.Frame, *cdp.Node, runtime.ExecutionContextID, bool) {
+	t.frameMu.RLock()
+	frame := t.frames[t.cur]
+	execCtx := t.execContexts[t.cur]
+	t.frameMu.RUnlock()
+
+	// the frame hasn't loaded yet.
+	if frame == nil || execCtx == 0 {
+		return nil, nil, 0, false
+	}
+
+	frame.RLock()
+	root := frame.Root
+	frame.RUnlock()
+
+	if root == nil {
+		// not root node yet?
+		return nil, nil, 0, false
+	}
+	return frame, root, execCtx, true
+}
+
 func (t *Target) run(ctx context.Context) {
 	type eventValue struct {
 		method cdproto.MethodType

testdata/poll.html

@@ -13,12 +13,6 @@
             }
         }
         setTimeout(() => globalThis.__FOO = 1, timeout);
-
-        if (hash === "mutation") {
-            setTimeout(() => {
-                document.body.appendChild(document.createElement('div'))
-            }, 100);
-        }
     </script>
 </head>
 <body>