diff --git a/query.go b/query.go index 840b4d0c..5d22c28a 100644 --- a/query.go +++ b/query.go @@ -26,12 +26,13 @@ type QueryAction Action // See Query for information on building an element selector and relevant // options. type Selector struct { - sel interface{} - fromNode *cdp.Node - exp int - by func(context.Context, *cdp.Node) ([]cdp.NodeID, error) - wait func(context.Context, *cdp.Frame, runtime.ExecutionContextID, ...cdp.NodeID) ([]*cdp.Node, error) - after func(context.Context, runtime.ExecutionContextID, ...*cdp.Node) error + sel interface{} + fromNode *cdp.Node + retryInterval time.Duration + exp int + by func(context.Context, *cdp.Node) ([]cdp.NodeID, error) + wait func(context.Context, *cdp.Frame, runtime.ExecutionContextID, ...cdp.NodeID) ([]*cdp.Node, error) + after func(context.Context, runtime.ExecutionContextID, ...*cdp.Node) error } // Query is a query action that queries the browser for specific element @@ -128,8 +129,9 @@ type Selector struct { // element nodes matching the selector. func Query(sel interface{}, opts ...QueryOption) QueryAction { s := &Selector{ - sel: sel, - exp: 1, + sel: sel, + exp: 1, + retryInterval: 5 * time.Millisecond, } // apply options @@ -155,7 +157,7 @@ func (s *Selector) Do(ctx context.Context) error { if t == nil { return ErrInvalidTarget } - return retryWithSleep(ctx, 5*time.Millisecond, func(ctx context.Context) (bool, error) { + return retryWithSleep(ctx, s.retryInterval, func(ctx context.Context) (bool, error) { frame, root, execCtx, ok := t.ensureFrame() if !ok { return false, nil @@ -553,6 +555,16 @@ func AtLeast(n int) QueryOption { } } +// RetryInterval is an element query action option to set the retry interval to specify +// how often it should retry when it failed to select the target element(s). +// +// The default value is 5ms. +func RetryInterval(interval time.Duration) QueryOption { + return func(s *Selector) { + s.retryInterval = interval + } +} + // After is an element query option that sets a func to execute after the // matched nodes have been returned by the browser, and after the node // condition is true. diff --git a/query_test.go b/query_test.go index 8801dddf..0bff1bf5 100644 --- a/query_test.go +++ b/query_test.go @@ -17,6 +17,7 @@ import ( "github.com/chromedp/cdproto/cdp" "github.com/chromedp/cdproto/css" "github.com/chromedp/cdproto/dom" + cdpruntime "github.com/chromedp/cdproto/runtime" "github.com/chromedp/chromedp/kb" ) @@ -188,6 +189,67 @@ func TestAtLeast(t *testing.T) { } } +func TestRetryInterval(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + opts []QueryOption + wantCountMin int + wantCountMax int + }{ + { + name: "default", + opts: []QueryOption{}, + // in 100ms + wantCountMin: 10, + wantCountMax: 20, + }, + { + name: "large interval", + opts: []QueryOption{RetryInterval(60 * time.Millisecond)}, + // in 100ms + wantCountMin: 2, + wantCountMax: 2, + }, + } + + ctx, cancel := testAllocate(t, "js.html") + defer cancel() + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + retryCount := 0 + + // count is a wait function that makes the query always fail and + // counts the number of retries. Note that the wait func is called + // only after the number of result nodes >= s.exp . + count := WaitFunc( + func(ctx context.Context, f *cdp.Frame, eci cdpruntime.ExecutionContextID, ni ...cdp.NodeID) ([]*cdp.Node, error) { + retryCount += 1 + return nil, ErrInvalidTarget + }, + ) + + ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancel() + + opts := append(tc.opts, count) + err := Run(ctx, Query("//input", opts...)) + + if err == nil || !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("want error context.DeadlineExceeded, got: %v", err) + } + if retryCount < tc.wantCountMin { + t.Fatalf("want retry count > %d, got: %d", tc.wantCountMin, retryCount) + } + if retryCount > tc.wantCountMax { + t.Fatalf("want retry count < %d, got: %d", tc.wantCountMax, retryCount) + } + }) + } +} + func TestByJSPath(t *testing.T) { t.Parallel()