add a ModifyCmdFunc ExecAllocatorOption

allocate.go

@@ -107,6 +107,8 @@ type ExecAllocator struct {
 	wg sync.WaitGroup
 
 	combinedOutputWriter io.Writer
+
+	modifyCmdFunc func(cmd *exec.Cmd)
 }
 
 // allocTempDir is used to group all ExecAllocator temporary user data dirs in
@@ -169,16 +171,24 @@ func (a *ExecAllocator) Allocate(ctx context.Context, opts ...BrowserOption) (*B
 			os.RemoveAll(dataDir)
 		}
 	}()
-	allocateCmdOptions(cmd)
+
+	if a.modifyCmdFunc != nil {
+		a.modifyCmdFunc(cmd)
+	} else {
+		allocateCmdOptions(cmd)
+	}
 
 	stdout, err := cmd.StdoutPipe()
 	if err != nil {
 		return nil, err
 	}
 	cmd.Stderr = cmd.Stdout
 
-	if len(a.initEnv) > 0 {
-		cmd.Env = append(os.Environ(), a.initEnv...)
+	// Preserve environment variables set in the (lowest priority) existing
+	// environment, OverrideCmdFunc(), and Env (highest priority)
+	if len(a.initEnv) > 0 || len(cmd.Env) > 0 {
+		cmd.Env = append(os.Environ(), cmd.Env...)
+		cmd.Env = append(cmd.Env, a.initEnv...)
 	}
 
 	// We must start the cmd before calling cmd.Wait, as otherwise the two
@@ -379,6 +389,16 @@ func Env(vars ...string) ExecAllocatorOption {
 	}
 }
 
+// ModifyCmdFunc allows for running an arbitrary function on the
+// browser exec.Cmd object. This overrides the default version
+// of the command which sends SIGKILL to any open browsers when
+// the Go program exits.
+func ModifyCmdFunc(f func(cmd *exec.Cmd)) ExecAllocatorOption {
+	return func(a *ExecAllocator) {
+		a.modifyCmdFunc = f
+	}
+}
+
 // UserDataDir is the command line option to set the user data dir.
 //
 // Note: set this option to manually set the profile directory used by Chrome.

allocate_test.go

@@ -346,3 +346,31 @@ func TestWithBrowserOptionAlreadyAllocated(t *testing.T) {
 		WithLogf(func(format string, args ...interface{}) {}),
 	)
 }
+
+func TestModifyCmdFunc(t *testing.T) {
+	t.Parallel()
+
+	tz := "Atlantic/Reykjavik"
+	allocCtx, cancel := NewExecAllocator(context.Background(),
+		append([]ExecAllocatorOption{
+			ModifyCmdFunc(func(cmd *exec.Cmd) {
+				cmd.Env = append(cmd.Env, "TZ=" + tz)
+			}),
+		}, allocOpts...)...)
+	defer cancel()
+
+	ctx, cancel := NewContext(allocCtx)
+	defer cancel()
+
+	var ret string
+	if err := Run(ctx,
+		Evaluate(`Intl.DateTimeFormat().resolvedOptions().timeZone`, &ret),
+	); err != nil {
+		t.Fatal(err)
+	}
+
+	if ret != tz {
+		t.Fatalf("got %s, want %s", ret, tz)
+	}
+}
+