From 97e80a00d57229ca619ec7cb9fdea2b954036904 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Sat, 6 Apr 2019 21:58:54 +0200
Subject: [PATCH] make NewAllocator's cancel func block on Wait

This way, the simple examples and tests don't need to do that
separately. Practically all users will want this cleanup work to be
synchronous, and practically all Go APIs are synchronous by default, so
this makes chromedp easier to use.
---
 allocate.go      | 14 +++++++++++---
 allocate_test.go | 23 ++++++++---------------
 chromedp_test.go |  1 -
 example_test.go  | 12 ------------
 input_test.go    | 10 ++++------
 5 files changed, 23 insertions(+), 37 deletions(-)

diff --git a/allocate.go b/allocate.go
index 3bae556..34d7060 100644
--- a/allocate.go
+++ b/allocate.go
@@ -22,8 +22,12 @@ type Allocator interface {
 	// as temporary directories) will be freed.
 	Allocate(context.Context) (*Browser, error)
 
-	// Wait can be called after cancelling an allocator's context, to block
-	// until all of its resources have been freed.
+	// TODO: Wait should probably return an error, which can then be
+	// retrieved by the user if just calling cancel().
+
+	// Wait blocks until an allocator has freed all of its resources.
+	// Cancelling the context obtained via NewAllocator will already perform
+	// this operation, so normally there's no need to call Wait directly.
 	Wait()
 }
 
@@ -38,7 +42,11 @@ func NewAllocator(parent context.Context, opts ...AllocatorOption) (context.Cont
 	}
 
 	ctx = context.WithValue(ctx, contextKey{}, c)
-	return ctx, cancel
+	cancelWait := func() {
+		cancel()
+		c.Allocator.Wait()
+	}
+	return ctx, cancelWait
 }
 
 // AllocatorOption is a allocator option.
diff --git a/allocate_test.go b/allocate_test.go
index 08f4496..e8d2b5b 100644
--- a/allocate_test.go
+++ b/allocate_test.go
@@ -30,16 +30,13 @@ func TestExecAllocator(t *testing.T) {
 		t.Fatalf("wanted %q, got %q", want, got)
 	}
 
+	// TODO: make cancel() work here too
+	poolCancel()
+
 	tempDir := FromContext(taskCtx).Browser.userDataDir
-	pool := FromContext(taskCtx).Allocator
-
-	cancel()
-	pool.Wait()
-
-	if _, err := os.Lstat(tempDir); os.IsNotExist(err) {
-		return
+	if _, err := os.Lstat(tempDir); !os.IsNotExist(err) {
+		t.Fatalf("temporary user data dir %q not deleted", tempDir)
 	}
-	t.Fatalf("temporary user data dir %q not deleted", tempDir)
 }
 
 func TestExecAllocatorCancelParent(t *testing.T) {
@@ -56,15 +53,11 @@ func TestExecAllocatorCancelParent(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	tempDir := FromContext(taskCtx).Browser.userDataDir
-	pool := FromContext(taskCtx).Allocator
-
 	// Canceling the pool context should stop all browsers too.
 	poolCancel()
-	pool.Wait()
 
-	if _, err := os.Lstat(tempDir); os.IsNotExist(err) {
-		return
+	tempDir := FromContext(taskCtx).Browser.userDataDir
+	if _, err := os.Lstat(tempDir); !os.IsNotExist(err) {
+		t.Fatalf("temporary user data dir %q not deleted", tempDir)
 	}
-	t.Fatalf("temporary user data dir %q not deleted", tempDir)
 }
diff --git a/chromedp_test.go b/chromedp_test.go
index b36a7b8..eec2d5d 100644
--- a/chromedp_test.go
+++ b/chromedp_test.go
@@ -64,6 +64,5 @@ func TestMain(m *testing.M) {
 	code := m.Run()
 
 	cancel()
-	FromContext(allocCtx).Allocator.Wait()
 	os.Exit(code)
 }
diff --git a/example_test.go b/example_test.go
index ce660a5..cab7315 100644
--- a/example_test.go
+++ b/example_test.go
@@ -26,10 +26,6 @@ func ExampleTitle() {
 
 	fmt.Println(title)
 
-	// wait for the resources to be cleaned up
-	cancel()
-	chromedp.FromContext(ctx).Allocator.Wait()
-
 	// no expected output, to not run this test as part of 'go test'; it's
 	// too slow, requiring internet access.
 }
@@ -69,10 +65,6 @@ func ExampleExecAllocator() {
 	lines := bytes.Split(bs, []byte("\n"))
 	fmt.Printf("DevToolsActivePort has %d lines\n", len(lines))
 
-	// wait for the resources to be cleaned up
-	cancel()
-	chromedp.FromContext(allocCtx).Allocator.Wait()
-
 	// Output:
 	// DevToolsActivePort has 2 lines
 }
@@ -101,10 +93,6 @@ func ExampleNewContext_manyTabs() {
 	fmt.Printf("Same browser: %t\n", c1.Browser == c2.Browser)
 	fmt.Printf("Same tab: %t\n", c1.Target == c2.Target)
 
-	// wait for the resources to be cleaned up
-	cancel()
-	c1.Allocator.Wait()
-
 	// Output:
 	// Same browser: true
 	// Same tab: false
diff --git a/input_test.go b/input_test.go
index 233a093..60c74b2 100644
--- a/input_test.go
+++ b/input_test.go
@@ -9,15 +9,13 @@ import (
 	"github.com/chromedp/cdproto/input"
 )
 
-const (
-	// inViewportJS is a javascript snippet that will get the specified node
-	// position relative to the viewport and returns true if the specified node
-	// is within the window's viewport.
-	inViewportJS = `(function(a) {
+// inViewportJS is a javascript snippet that will get the specified node
+// position relative to the viewport and returns true if the specified node
+// is within the window's viewport.
+const inViewportJS = `(function(a) {
 		var r = a[0].getBoundingClientRect();
 		return r.top >= 0 && r.left >= 0 && r.bottom <= window.innerHeight && r.right <= window.innerWidth;
 	})($x('%s'))`
-)
 
 func TestMouseClickXY(t *testing.T) {
 	t.Parallel()