testing: add Context

Adds a new Context method to testing.T, that returns a context, that is
canceled before the end of its test function.

Fixes golang#36532.

src/testing/testing.go

@@ -371,6 +371,7 @@ package testing
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"flag"
 	"fmt"
@@ -632,6 +633,9 @@ type common struct {
 	tempDir    string
 	tempDirErr error
 	tempDirSeq int32
+
+	ctx       context.Context
+	cancelCtx context.CancelFunc
 }
 
 // Short reports whether the -test.short flag is set.
@@ -896,6 +900,7 @@ type TB interface {
 	Skipf(format string, args ...any)
 	Skipped() bool
 	TempDir() string
+	Context() context.Context
 
 	// A private method to prevent users implementing the
 	// interface and so future additions to it will not
@@ -1307,6 +1312,16 @@ func (c *common) Setenv(key, value string) {
 	}
 }
 
+// Context returns a context that is cancelled just before
+// [T.Cleanup]-registered functions are called.
+//
+// Cleanup functions can wait for any resources
+// that shut down on Context.Done before the test completes.
+func (c *common) Context() context.Context {
+	c.checkFuzzFn("Context")
+	return c.ctx
+}
+
 // panicHandling controls the panic handling used by runCleanup.
 type panicHandling int
 
@@ -1339,6 +1354,10 @@ func (c *common) runCleanup(ph panicHandling) (panicVal any) {
 		}
 	}()
 
+	if c.cancelCtx != nil {
+		c.cancelCtx()
+	}
+
 	for {
 		var cleanup func()
 		c.mu.Lock()
@@ -1716,15 +1735,23 @@ func (t *T) Run(name string, f func(t *T)) bool {
 	// continue walking the stack into the parent test.
 	var pc [maxStackLen]uintptr
 	n := runtime.Callers(2, pc[:])
+
+	parentCtx := t.Context()
+	if parentCtx == nil {
+		parentCtx = context.Background()
+	}
+	ctx, cancelCtx := context.WithCancel(parentCtx)
 	t = &T{
 		common: common{
-			barrier: make(chan bool),
-			signal:  make(chan bool, 1),
-			name:    testName,
-			parent:  &t.common,
-			level:   t.level + 1,
-			creator: pc[:n],
-			chatty:  t.chatty,
+			barrier:   make(chan bool),
+			signal:    make(chan bool, 1),
+			name:      testName,
+			parent:    &t.common,
+			level:     t.level + 1,
+			creator:   pc[:n],
+			chatty:    t.chatty,
+			ctx:       ctx,
+			cancelCtx: cancelCtx,
 		},
 		context: t.context,
 	}
@@ -2150,15 +2177,18 @@ func runTests(matchString func(pat, str string) (bool, error), tests []InternalT
 				// to keep trying.
 				break
 			}
-			ctx := newTestContext(*parallel, newMatcher(matchString, *match, "-test.run", *skip))
-			ctx.deadline = deadline
+			ctx, cancelCtx := context.WithCancel(context.Background())
+			tc := newTestContext(*parallel, newMatcher(matchString, *match, "-test.run", *skip))
+			tc.deadline = deadline
 			t := &T{
 				common: common{
-					signal:  make(chan bool, 1),
-					barrier: make(chan bool),
-					w:       os.Stdout,
+					signal:    make(chan bool, 1),
+					barrier:   make(chan bool),
+					w:         os.Stdout,
+					ctx:       ctx,
+					cancelCtx: cancelCtx,
 				},
-				context: ctx,
+				context: tc,
 			}
 			if Verbose() {
 				t.chatty = newChattyPrinter(t.w)

src/testing/testing_test.go

@@ -6,6 +6,8 @@ package testing_test
 
 import (
 	"bytes"
+	"context"
+	"errors"
 	"fmt"
 	"internal/race"
 	"internal/testenv"
@@ -813,3 +815,29 @@ func TestParentRun(t1 *testing.T) {
 		})
 	})
 }
+
+func TestContext(t *testing.T) {
+	ctx := t.Context()
+	if err := ctx.Err(); err != nil {
+		t.Fatalf("expected non-canceled context, got %v", err)
+	}
+
+	var innerCtx context.Context
+	t.Run("inner", func(t *testing.T) {
+		innerCtx = t.Context()
+		if err := innerCtx.Err(); err != nil {
+			t.Fatalf("expected inner test to not inherit canceled context, got %v", err)
+		}
+	})
+	t.Run("inner2", func(t *testing.T) {
+		if !errors.Is(innerCtx.Err(), context.Canceled) {
+			t.Fatal("expected context of sibling test to be canceled after its test function finished")
+		}
+	})
+
+	t.Cleanup(func() {
+		if !errors.Is(ctx.Err(), context.Canceled) {
+			t.Fatal("expected context canceled before cleanup")
+		}
+	})
+}