From d7fa4ab94d0f288aae4dc004337ce9053649cb7f Mon Sep 17 00:00:00 2001
From: Boris Nagaev <bnagaev@gmail.com>
Date: Wed, 31 Jul 2024 14:07:51 -0300
Subject: [PATCH] test.Guard: make timeout configurable

---
 test/timeout.go | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/test/timeout.go b/test/timeout.go
index 52d4762..4370ced 100644
--- a/test/timeout.go
+++ b/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)