test cancelling an entire Allocator directly

To ensure that it propagates to each browser correctly.

allocate_test.go

@@ -9,8 +9,8 @@ import (
 func TestExecAllocator(t *testing.T) {
 	t.Parallel()
 
-	poolCtx, cancel := NewAllocator(context.Background(), WithExecAllocator(allocOpts...))
-	defer cancel()
+	poolCtx, poolCancel := NewAllocator(context.Background(), WithExecAllocator(allocOpts...))
+	defer poolCancel()
 
 	// TODO: test that multiple child contexts are run in different
 	// processes and browsers.
@@ -41,3 +41,30 @@ func TestExecAllocator(t *testing.T) {
 	}
 	t.Fatalf("temporary user data dir %q not deleted", tempDir)
 }
+
+func TestExecAllocatorCancelParent(t *testing.T) {
+	t.Parallel()
+
+	poolCtx, poolCancel := NewAllocator(context.Background(), WithExecAllocator(allocOpts...))
+	defer poolCancel()
+
+	// TODO: test that multiple child contexts are run in different
+	// processes and browsers.
+
+	taskCtx, _ := NewContext(poolCtx)
+	if err := Run(taskCtx, Tasks{}); err != nil {
+		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
+	}
+	t.Fatalf("temporary user data dir %q not deleted", tempDir)
+}