From 2f7606f1a34941061322fd0a583bf195136fe224 Mon Sep 17 00:00:00 2001
From: Adam Luzsi <adamluzsi@gmail.com>
Date: Tue, 4 Jun 2024 20:36:11 +0200
Subject: [PATCH] testcase.T#Done now allows you to be notified about when a
 test has completed

This makes it easy to be notified about the end of a test when you are working with goroutines and channels.
---
 T.go      | 34 +++++++++++++++++++++++-----------
 T_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 11 deletions(-)

diff --git a/T.go b/T.go
index bffe08e..dd6fe9e 100644
--- a/T.go
+++ b/T.go
@@ -2,7 +2,6 @@ package testcase
 
 import (
 	"math/rand"
-	"sync"
 	"testing"
 	"time"
 
@@ -35,9 +34,11 @@ func newT(tb testing.TB, spec *Spec) *T {
 		Random: random.New(rand.NewSource(spec.getTestSeed(tb))),
 		It:     assert.MakeIt(tb),
 
-		spec:     spec,
+		spec: spec,
+		tags: spec.getTagSet(),
+
 		vars:     newVariables(),
-		tags:     spec.getTagSet(),
+		done:     make(chan struct{}),
 		teardown: &teardown.Teardown{CallerOffset: 1},
 	}
 }
@@ -63,16 +64,14 @@ type T struct {
 	// but mark test failed on a failed assertion.
 	assert.It
 
-	spec     *Spec
+	spec *Spec
+	tags map[string]struct{}
+
 	vars     *variables
-	tags     map[string]struct{}
+	done     chan struct{}
 	teardown *teardown.Teardown
 
-	depsInit sync.Once
-	deps     map[string]struct{}
-
-	// TODO: protect it against concurrency
-	timerPaused bool
+	timerPaused bool // TODO: protect it against concurrency
 
 	cache struct {
 		contexts []*Spec
@@ -116,6 +115,9 @@ func (t *T) setUp() func() {
 	t.TB.Helper()
 	t.vars.reset()
 
+	done := make(chan struct{})
+	t.done = done
+
 	contexts := t.contexts()
 	for _, c := range contexts {
 		t.vars.merge(c.vars)
@@ -133,7 +135,10 @@ func (t *T) setUp() func() {
 		}
 	}
 
-	return t.teardown.Finish
+	return func() {
+		t.teardown.Finish()
+		close(done)
+	}
 }
 
 func (t *T) HasTag(tag string) bool {
@@ -271,3 +276,10 @@ func (t *T) LogPretty(vs ...any) {
 	}
 	t.Log(args...)
 }
+
+// Done function notifies the end of the test.
+// If a test involves goroutines, listening to the done channel from the test
+// can notify them about the test's end, preventing goroutine leaks.
+func (t *T) Done() <-chan struct{} {
+	return t.done
+}
diff --git a/T_test.go b/T_test.go
index 0c0e4cc..7cf8393 100644
--- a/T_test.go
+++ b/T_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math/rand"
 	"os"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -599,3 +600,54 @@ func TestT_LogPretty(t *testing.T) {
 	assert.Contain(t, dtb.Logs.String(), "[]int{\n\t1,\n\t2,\n\t4,\n}")
 	assert.Contain(t, dtb.Logs.String(), "testcase_test.X{\n\tFoo: \"hello\",\n}")
 }
+
+func ExampleT_Done() {
+	s := testcase.NewSpec(nil)
+
+	s.Test("", func(t *testcase.T) {
+		go func() {
+			select {
+			// case do something for the test
+			case <-t.Done():
+				return // test is over, time to garbage collect
+			}
+		}()
+	})
+}
+
+func TestT_Done(t *testing.T) {
+	s := testcase.NewSpec(t)
+
+	var isdone = func(t *testcase.T) bool {
+		select {
+		case <-t.Done():
+			return true
+		default:
+			return false
+		}
+	}
+
+	var done int32
+	s.Test("", func(t *testcase.T) {
+		assert.False(t, isdone(t))
+		go func() {
+			<-t.Done() // after the test is done
+			atomic.AddInt32(&done, 1)
+		}()
+		t.Cleanup(func() {
+			assert.False(t, isdone(t),
+				"during cleanup the done should be not ready")
+
+			t.Cleanup(func() {
+				assert.False(t, isdone(t),
+					"during a cleanup of cleanup, done should not be ready")
+			})
+		})
+	})
+
+	s.Finish()
+
+	assert.Eventually(t, time.Second, func(t assert.It) {
+		assert.Equal(t, atomic.LoadInt32(&done), 1)
+	})
+}