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()