add Page.WaitPauseOpen

related to go-rod/stealth#2
go-rod · Aug 4, 2020 · 85ac15b · 85ac15b
1 parent 4391538
commit 85ac15b
Show file tree

Hide file tree

Showing 5 changed files with 83 additions and 4 deletions.
diff --git a/fixtures/open-page-subpage.html b/fixtures/open-page-subpage.html
@@ -0,0 +1,9 @@
+<html>
+    <body>
+        <script>
+            if (!window.a) {
+                window.a = 'new page'
+            }
+        </script>
+    </body>
+</html>
diff --git a/fixtures/open-page.html b/fixtures/open-page.html
@@ -1,5 +1,5 @@
 <html>
     <body>
-        <a href="./click.html" id="link" target="_blank">open page</a>
+        <a href="./open-page-subpage.html" id="link" target="_blank">open page</a>
     </body>
 </html>
diff --git a/page.go b/page.go
@@ -215,7 +215,8 @@ func (p *Page) CloseE() error {
 	return nil
 }
 
-// HandleDialogE doc is similar to the method HandleDialog
+// HandleDialogE doc is similar to the method HandleDialog.
+// Because the js will be paused, you should put the code that triggers in a goroutine.
 func (p *Page) HandleDialogE(accept bool, promptText string) func() error {
 	recover := p.EnableDomain(&proto.PageEnable{})
 
@@ -301,6 +302,35 @@ func (p *Page) WaitOpenE() func() (*Page, error) {
 	}
 }
 
+// WaitPauseOpenE waits for a page opened by the current page, before opening pause the js execution.
+// Because the js will be paused, you should put the code that triggers it in a goroutine.
+func (p *Page) WaitPauseOpenE() (wait func() (*Page, error), resume func() error, err error) {
+	wait = p.WaitOpenE()
+
+	// TODO: we have to use the browser to call, seems like a chrome bug
+	err = proto.TargetSetAutoAttach{
+		AutoAttach:             true,
+		WaitForDebuggerOnStart: true,
+		Flatten:                true,
+	}.Call(p.browser.Context(p.ctx, p.ctxCancel))
+	if err != nil {
+		return
+	}
+
+	resume = func() error {
+		err = proto.TargetSetAutoAttach{
+			Flatten: true,
+		}.Call(p.browser.Context(p.ctx, p.ctxCancel))
+		if err != nil {
+			return err
+		}
+
+		return proto.RuntimeRunIfWaitingForDebugger{}.Call(p)
+	}
+
+	return
+}
+
 // PauseE doc is similar to the method Pause
 func (p *Page) PauseE() error {
 	wait := p.WaitEvent(&proto.DebuggerResumed{})

diff --git a/page_test.go b/page_test.go
@@ -232,7 +232,7 @@ func (s *S) TestPageExposeJSHelper() {
 	s.Equal("object", page.Eval("typeof(rod)").Str)
 }
 
-func (s *S) TestUntilPage() {
+func (s *S) TestPageWaitOpen() {
 	page := s.page.Timeout(3 * time.Second).Navigate(srcFile("fixtures/open-page.html"))
 	defer page.CancelTimeout()
 
@@ -241,8 +241,35 @@ func (s *S) TestUntilPage() {
 	page.Element("a").Click()
 
 	newPage := wait()
+	defer newPage.Close()
 
-	s.Equal("click me", newPage.Element("button").Text())
+	s.Equal("new page", newPage.Eval("window.a").String())
+}
+
+func (s *S) TestPageWaitPauseOpen() {
+	page := s.page.Timeout(3 * time.Second).Navigate(srcFile("fixtures/open-page.html"))
+	defer page.CancelTimeout()
+
+	wait, resume := page.WaitPauseOpen()
+
+	go page.Element("a").Click()
+
+	newPage := wait()
+
+	newPage.EvalOnNewDocument(`window.a = 'ok'`)
+	defer newPage.Close()
+	resume()
+
+	s.Equal("ok", newPage.Eval(`window.a`).String())
+
+	w := page.WaitOpen()
+
+	page.Element("a").Click()
+
+	newPage = w()
+	defer newPage.Close()
+
+	s.Equal("new page", newPage.Eval("window.a").String())
 }
 
 func (s *S) TestPageWait() {

diff --git a/sugar.go b/sugar.go
@@ -235,6 +235,19 @@ func (p *Page) WaitOpen() (wait func() (newPage *Page)) {
 	}
 }
 
+// WaitPauseOpen waits for a page opened by the current page, before opening pause the js execution.
+// Because the js will be paused, you should put the code that triggers it in a goroutine, such as the click.
+func (p *Page) WaitPauseOpen() (wait func() *Page, resume func()) {
+	newPage, r, err := p.WaitPauseOpenE()
+	utils.E(err)
+
+	return func() *Page {
+		page, err := newPage()
+		utils.E(err)
+		return page
+	}, func() { utils.E(r()) }
+}
+
 // Pause stops on the next JavaScript statement
 func (p *Page) Pause() *Page {
 	utils.E(p.PauseE())