feat(testutil): make utils t.Context() aware

mafredri · mafredri · commit fa4063c87e50 · 2025-04-28T11:48:19.000Z
diff --git a/testutil/chan.go b/testutil/chan.go
@@ -10,36 +10,37 @@ import (
 // the channel is closed, the zero value of the channel type will be returned.
 //
 // Safety: Must only be called from the Go routine that created `t`.
-func TryReceive[A any](ctx context.Context, t testing.TB, c <-chan A) A {
+func TryReceive[A any](ctx context.Context, t testing.TB, c <-chan A) (a A) {
 	t.Helper()
 	select {
+	case <-t.Context().Done():
+		t.Fatal("test timeout")
 	case <-ctx.Done():
-		t.Fatal("timeout")
-		var a A
-		return a
-	case a := <-c:
-		return a
+		t.Fatal("context timeout")
+	case a = <-c:
 	}
+	return a
 }
 
 // RequireReceive will receive a value from the chan and return it. If the
 // context expires or the channel is closed before a value can be received,
 // it will fail the test.
 //
 // Safety: Must only be called from the Go routine that created `t`.
-func RequireReceive[A any](ctx context.Context, t testing.TB, c <-chan A) A {
+func RequireReceive[A any](ctx context.Context, t testing.TB, c <-chan A) (a A) {
 	t.Helper()
+	var ok bool
 	select {
+	case <-t.Context().Done():
+		t.Fatal("test timeout")
 	case <-ctx.Done():
-		t.Fatal("timeout")
-		var a A
-		return a
-	case a, ok := <-c:
+		t.Fatal("context timeout")
+	case a, ok = <-c:
 		if !ok {
 			t.Fatal("channel closed")
 		}
-		return a
 	}
+	return a
 }
 
 // RequireSend will send the given value over the chan and then return. If
@@ -49,8 +50,10 @@ func RequireReceive[A any](ctx context.Context, t testing.TB, c <-chan A) A {
 func RequireSend[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
 	t.Helper()
 	select {
+	case <-t.Context().Done():
+		t.Fatal("test timeout")
 	case <-ctx.Done():
-		t.Fatal("timeout")
+		t.Fatal("context timeout")
 	case c <- a:
 		// OK!
 	}
diff --git a/testutil/ctx.go b/testutil/ctx.go
@@ -7,7 +7,7 @@ import (
 )
 
 func Context(t *testing.T, dur time.Duration) context.Context {
-	ctx, cancel := context.WithTimeout(context.Background(), dur)
+	ctx, cancel := context.WithTimeout(t.Context(), dur)
 	t.Cleanup(cancel)
 	return ctx
 }
diff --git a/testutil/eventually.go b/testutil/eventually.go
@@ -41,6 +41,9 @@ func Eventually(ctx context.Context, t testing.TB, condition func(ctx context.Co
 	defer ticker.Stop()
 	for tick := ticker.C; ; {
 		select {
+		case <-t.Context().Done():
+			assert.NoError(t, t.Context().Err(), msg)
+			return false
 		case <-ctx.Done():
 			assert.NoError(t, ctx.Err(), msg)
 			return false
diff --git a/testutil/eventually_test.go b/testutil/eventually_test.go
@@ -25,7 +25,7 @@ func TestEventually(t *testing.T) {
 			}()
 			return state > 2
 		}
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+		ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitShort)
 		defer cancel()
 		testutil.Eventually(ctx, t, condition, testutil.IntervalFast)
 	})
@@ -36,7 +36,7 @@ func TestEventually(t *testing.T) {
 		panicky := func() {
 			mockT := new(testing.T)
 			condition := func(_ context.Context) bool { return true }
-			testutil.Eventually(context.Background(), mockT, condition, testutil.IntervalFast)
+			testutil.Eventually(t.Context(), mockT, condition, testutil.IntervalFast)
 		}
 		assert.Panics(t, panicky)
 	})
diff --git a/testutil/pty.go b/testutil/pty.go
@@ -83,6 +83,8 @@ func (tr *TerminalReader) ReadUntil(ctx context.Context, matcher func(line strin
 			if err != nil {
 				return err
 			}
+		case <-tr.t.Context().Done():
+			return tr.t.Context().Err()
 		case <-ctx.Done():
 			return ctx.Err()
 		}

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ import (`
`7`	`7`	`)`
`8`	`8`
`9`	`9`	`func Context(t *testing.T, dur time.Duration) context.Context {`
`10`		`- ctx, cancel := context.WithTimeout(context.Background(), dur)`
	`10`	`+ ctx, cancel := context.WithTimeout(t.Context(), dur)`
`11`	`11`	`t.Cleanup(cancel)`
`12`	`12`	`return ctx`
`13`	`13`	`}`
Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,8 @@ func (tr *TerminalReader) ReadUntil(ctx context.Context, matcher func(line strin`
`83`	`83`	`if err != nil {`
`84`	`84`	`return err`
`85`	`85`	`}`
	`86`	`+ case <-tr.t.Context().Done():`
	`87`	`+ return tr.t.Context().Err()`
`86`	`88`	`case <-ctx.Done():`
`87`	`89`	`return ctx.Err()`
`88`	`90`	`}`