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)