testing: support B.Context and F.Context

CL 603959 added T.Context for #36532.

The discussion on the proposal only mentions t.Context.
However, the implementation of CL 603959 also added B.Context and F.Context.
They were added to the API listing, and B.Context was mentioned in
the release notes.

Unfortunately, the new B.Context and F.Context methods always
returned nil, rather than a context.Context value.

This change adds a working implementation of B.Context and F.Context.

For #36532
Fixes #70866

Change-Id: I8a44e6649fb658e4f641ffb7efd08b4374f578ef
Reviewed-on: https://go-review.googlesource.com/c/go/+/637236
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Chase <drchase@google.com>
Reviewed-by: Damien Neil <dneil@google.com>

src/cmd/go/testdata/script/test_fuzz_context.txt

@@ -0,0 +1,47 @@
+[!fuzz] skip
+[short] skip
+env GOCACHE=$WORK/cache
+
+# Test fuzz.Context.
+go test -vet=off context_fuzz_test.go
+stdout ^ok
+! stdout FAIL
+
+go test -vet=off -fuzz=Fuzz -fuzztime=1x context_fuzz_test.go
+stdout ok
+! stdout FAIL
+
+-- context_fuzz_test.go --
+package context_fuzz
+
+import (
+	"context"
+	"errors"
+	"testing"
+)
+
+func Fuzz(f *testing.F) {
+	ctx := f.Context()
+	if err := ctx.Err(); err != nil {
+		f.Fatalf("expected non-canceled context, got %v", err)
+	}
+
+	f.Fuzz(func(t *testing.T, data []byte) {
+		innerCtx := t.Context()
+		if err := innerCtx.Err(); err != nil {
+			t.Fatalf("expected inner test to not inherit canceled context, got %v", err)
+		}
+
+		t.Cleanup(func() {
+			if !errors.Is(innerCtx.Err(), context.Canceled) {
+				t.Fatal("expected context of inner test to be canceled after its fuzz function finished")
+			}
+		})
+	})
+
+	f.Cleanup(func() {
+		if !errors.Is(ctx.Err(), context.Canceled) {
+			f.Fatal("expected context canceled before cleanup")
+		}
+	})
+}

src/testing/benchmark.go

@@ -5,6 +5,7 @@
 package testing
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"internal/sysinfo"
@@ -181,6 +182,7 @@ func (b *B) ReportAllocs() {
 func (b *B) runN(n int) {
 	benchmarkLock.Lock()
 	defer benchmarkLock.Unlock()
+	ctx, cancelCtx := context.WithCancel(context.Background())
 	defer func() {
 		b.runCleanup(normalPanic)
 		b.checkRaces()
@@ -191,6 +193,8 @@ func (b *B) runN(n int) {
 	b.resetRaces()
 	b.N = n
 	b.loopN = 0
+	b.ctx = ctx
+	b.cancelCtx = cancelCtx
 
 	b.parallelism = 1
 	b.ResetTimer()

src/testing/benchmark_test.go

@@ -7,6 +7,8 @@ package testing_test
 import (
 	"bytes"
 	"cmp"
+	"context"
+	"errors"
 	"runtime"
 	"slices"
 	"strings"
@@ -127,6 +129,34 @@ func TestRunParallelSkipNow(t *testing.T) {
 	})
 }
 
+func TestBenchmarkContext(t *testing.T) {
+	testing.Benchmark(func(b *testing.B) {
+		ctx := b.Context()
+		if err := ctx.Err(); err != nil {
+			b.Fatalf("expected non-canceled context, got %v", err)
+		}
+
+		var innerCtx context.Context
+		b.Run("inner", func(b *testing.B) {
+			innerCtx = b.Context()
+			if err := innerCtx.Err(); err != nil {
+				b.Fatalf("expected inner benchmark to not inherit canceled context, got %v", err)
+			}
+		})
+		b.Run("inner2", func(b *testing.B) {
+			if !errors.Is(innerCtx.Err(), context.Canceled) {
+				t.Fatal("expected context of sibling benchmark 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")
+			}
+		})
+	})
+}
+
 func ExampleB_RunParallel() {
 	// Parallel benchmark for text/template.Template.Execute on a single object.
 	testing.Benchmark(func(b *testing.B) {

src/testing/fuzz.go

@@ -5,6 +5,7 @@
 package testing
 
 import (
+	"context"
 	"errors"
 	"flag"
 	"fmt"
@@ -293,6 +294,8 @@ func (f *F) Fuzz(ff any) {
 			f.tstate.match.clearSubNames()
 		}
 
+		ctx, cancelCtx := context.WithCancel(f.ctx)
+
 		// Record the stack trace at the point of this call so that if the subtest
 		// function - which runs in a separate stack - is marked as a helper, we can
 		// continue walking the stack into the parent test.
@@ -307,6 +310,8 @@ func (f *F) Fuzz(ff any) {
 				level:     f.level + 1,
 				creator:   pc[:n],
 				chatty:    f.chatty,
+				ctx:       ctx,
+				cancelCtx: cancelCtx,
 			},
 			tstate: f.tstate,
 		}
@@ -508,6 +513,7 @@ func runFuzzTests(deps testDeps, fuzzTests []InternalFuzzTarget, deadline time.T
 						continue
 					}
 				}
+				ctx, cancelCtx := context.WithCancel(context.Background())
 				f := &F{
 					common: common{
 						signal:    make(chan bool),
@@ -516,6 +522,8 @@ func runFuzzTests(deps testDeps, fuzzTests []InternalFuzzTarget, deadline time.T
 						parent:    &root,
 						level:     root.level + 1,
 						chatty:    root.chatty,
+						ctx:       ctx,
+						cancelCtx: cancelCtx,
 					},
 					tstate: tstate,
 					fstate: fstate,
@@ -590,6 +598,7 @@ func runFuzzing(deps testDeps, fuzzTests []InternalFuzzTarget) (ok bool) {
 		return false
 	}
 
+	ctx, cancelCtx := context.WithCancel(context.Background())
 	f := &F{
 		common: common{
 			signal:    make(chan bool),
@@ -598,6 +607,8 @@ func runFuzzing(deps testDeps, fuzzTests []InternalFuzzTarget) (ok bool) {
 			parent:    &root,
 			level:     root.level + 1,
 			chatty:    root.chatty,
+			ctx:       ctx,
+			cancelCtx: cancelCtx,
 		},
 		fstate: fstate,
 		tstate: tstate,

src/testing/testing.go

@@ -1385,10 +1385,10 @@ func (c *common) Chdir(dir string) {
 }
 
 // Context returns a context that is canceled just before
-// [T.Cleanup]-registered functions are called.
+// Cleanup-registered functions are called.
 //
 // Cleanup functions can wait for any resources
-// that shut down on Context.Done before the test completes.
+// that shut down on Context.Done before the test or benchmark completes.
 func (c *common) Context() context.Context {
 	c.checkFuzzFn("Context")
 	return c.ctx