test.Guard: make timeout configurable

Author: starius

test/timeout.go

@@ -9,12 +9,34 @@ import (
 	"github.com/fortytw2/leaktest"
 )
 
+// GuardConfig stores options for Guard function.
+type GuardConfig struct {
+	timeout time.Duration
+}
+
+// GuardOption is an option for Guard function.
+type GuardOption func(*GuardConfig)
+
+// WithGuardTimeout sets timeout for the guard. Default is 5s.
+func WithGuardTimeout(timeout time.Duration) GuardOption {
+	return func(c *GuardConfig) {
+		c.timeout = timeout
+	}
+}
+
 // Guard implements a test level timeout.
-func Guard(t *testing.T) func() {
+func Guard(t *testing.T, opts ...GuardOption) func() {
+	cfg := GuardConfig{
+		timeout: 5 * time.Second,
+	}
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+
 	done := make(chan struct{})
 	go func() {
 		select {
-		case <-time.After(5 * time.Second):
+		case <-time.After(cfg.timeout):
 			err := pprof.Lookup("goroutine").WriteTo(os.Stdout, 1)
 			if err != nil {
 				panic(err)