From 8ff2971fc50db58525c8cffcbab6bea95e41fac8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Fri, 22 Mar 2019 17:19:30 +0000
Subject: [PATCH] test cancelling an entire Allocator directly

To ensure that it propagates to each browser correctly.
---
 allocate_test.go | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/allocate_test.go b/allocate_test.go
index 47d68bc..f184f50 100644
--- a/allocate_test.go
+++ b/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)
+}