From 6f22d3d6291477f099fa44fa2f453362806df23a Mon Sep 17 00:00:00 2001
From: Marcin Wyszynski <marcin.pixie@gmail.com>
Date: Thu, 22 Dec 2022 11:58:24 +0100
Subject: [PATCH] Better timeout handling

---
 goblin.go      | 32 +++++++++++++++++++++-----------
 goblin_test.go |  2 +-
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/goblin.go b/goblin.go
index a4617db..4c75986 100644
--- a/goblin.go
+++ b/goblin.go
@@ -3,6 +3,7 @@ package goblin
 import (
 	"flag"
 	"fmt"
+	"os"
 	"regexp"
 	"runtime"
 	"sync"
@@ -44,7 +45,7 @@ func (g *G) Describe(name string, h func()) {
 }
 
 func (g *G) Timeout(time time.Duration) {
-	g.timeout = time
+	g.currentTimeout = time
 	g.timer.Reset(time)
 }
 
@@ -130,7 +131,6 @@ type It struct {
 	failure   *Failure
 	failureMu sync.RWMutex
 	reporter  Reporter
-	isAsync   bool
 }
 
 func (it *It) run(g *G) bool {
@@ -171,7 +171,6 @@ type Xit struct {
 	parent   *Describe
 	failure  *Failure
 	reporter Reporter
-	isAsync  bool
 }
 
 func (xit *Xit) run(g *G) bool {
@@ -193,6 +192,13 @@ func parseFlags() {
 	} else {
 		runRegex = nil
 	}
+
+	if envTimeout := os.Getenv("GOBLIN_TIMEOUT"); envTimeout != "" {
+		var err error
+		if *timeout, err = time.ParseDuration(envTimeout); err != nil {
+			panic(err)
+		}
+	}
 }
 
 var doParseOnce sync.Once
@@ -202,11 +208,9 @@ var regexParam = flag.String("goblin.run", "", "Runs only tests which match the
 var runRegex *regexp.Regexp
 
 func Goblin(t *testing.T, arguments ...string) *G {
-	doParseOnce.Do(func() {
-		parseFlags()
-	})
+	doParseOnce.Do(parseFlags)
 
-	g := &G{t: t, timeout: *timeout}
+	g := &G{t: t, defaultTimeout: *timeout, currentTimeout: *timeout}
 	var fancy TextFancier
 	if *isTty {
 		fancy = &TerminalFancier{}
@@ -222,7 +226,7 @@ func runIt(g *G, it *It) {
 	g.mutex.Lock()
 	g.timedOut = false
 	g.mutex.Unlock()
-	g.timer = time.NewTimer(g.timeout)
+	g.timer = time.NewTimer(g.currentTimeout)
 	g.shouldContinue = make(chan bool)
 	if call, ok := it.h.(func()); ok {
 		// the test is synchronous
@@ -262,17 +266,18 @@ func runIt(g *G, it *It) {
 		//Set to nil as it shouldn't continue
 		g.shouldContinue = nil
 		g.timedOut = true
-		g.Fail("Test exceeded " + fmt.Sprintf("%s", g.timeout))
+		g.Fail("Test exceeded " + g.currentTimeout.String())
 	}
 	// Reset timeout value
-	g.timeout = *timeout
+	g.currentTimeout = g.defaultTimeout
 }
 
 type G struct {
 	t              *testing.T
 	parent         *Describe
 	currentIt      Itable
-	timeout        time.Duration
+	currentTimeout time.Duration
+	defaultTimeout time.Duration
 	reporter       Reporter
 	timedOut       bool
 	shouldContinue chan bool
@@ -280,6 +285,11 @@ type G struct {
 	timer          *time.Timer
 }
 
+func (g *G) SetDefaultTimeout(timeout time.Duration) {
+	g.defaultTimeout = timeout
+	g.currentTimeout = timeout
+}
+
 func (g *G) SetReporter(r Reporter) {
 	g.reporter = r
 }
diff --git a/goblin_test.go b/goblin_test.go
index 2d7ae81..5b949da 100644
--- a/goblin_test.go
+++ b/goblin_test.go
@@ -494,7 +494,7 @@ func TestItTimeout(t *testing.T) {
 		})
 
 		g.It("Should revert for different it", func() {
-			g.Assert(g.timeout).Equal(10 * time.Millisecond)
+			g.Assert(g.currentTimeout).Equal(10 * time.Millisecond)
 		})
 
 	})