Add ElementFunc as a RaceContext condition #534

query.go

@@ -386,6 +386,14 @@ func (rc *RaceContext) Element(selector string) *RaceContext {
 	return rc
 }
 
+// ElementFunc takes a custom function to determine race success
+func (rc *RaceContext) ElementFunc(fn func(*Page) (*Element, error)) *RaceContext {
+	rc.branches = append(rc.branches, &raceBranch{
+		condition: fn,
+	})
+	return rc
+}
+
 // ElementX the doc is similar to ElementX
 func (rc *RaceContext) ElementX(selector string) *RaceContext {
 	rc.branches = append(rc.branches, &raceBranch{

query_test.go

@@ -210,6 +210,14 @@ func TestPageRace(t *testing.T) {
 		MustHandle(func(e *rod.Element) { g.Eq("01", e.MustText()) }).MustDo()
 	g.Eq("01", p.Race().MustElementByJS("() => document.querySelector('button')", nil).MustDo().MustText())
 
+	raceFunc := func(p *rod.Page) (*rod.Element, error) {
+		el := p.MustElement("button")
+		g.Eq("01", el.MustText())
+		return el, nil
+	}
+	p.Race().ElementFunc(raceFunc).MustHandle(func(e *rod.Element) { g.Eq("01", e.MustText()) }).MustDo()
+	g.Eq("01", p.Race().ElementFunc(raceFunc).MustDo().MustText())
+
 	el, err := p.Sleeper(func() utils.Sleeper { return utils.CountSleeper(2) }).Race().
 		Element("not-exists").MustHandle(func(e *rod.Element) {}).
 		ElementX("//not-exists").