chore: update testutil chan helpers (#17408)

testutil/chan.go

@@ -0,0 +1,57 @@
+package testutil
+
+import (
+	"context"
+	"testing"
+)
+
+// TryReceive will attempt to receive a value from the chan and return it. If
+// the context expires before a value can be received, it will fail the test. If
+// 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 {
+	t.Helper()
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout")
+		var a A
+		return a
+	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 {
+	t.Helper()
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout")
+		var a A
+		return a
+	case a, ok := <-c:
+		if !ok {
+			t.Fatal("channel closed")
+		}
+		return a
+	}
+}
+
+// RequireSend will send the given value over the chan and then return. If
+// the context expires before the send succeeds, it will fail the test.
+//
+// Safety: Must only be called from the Go routine that created `t`.
+func RequireSend[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
+	t.Helper()
+	select {
+	case <-ctx.Done():
+		t.Fatal("timeout")
+	case c <- a:
+		// OK!
+	}
+}

testutil/ctx.go

@@ -11,37 +11,3 @@ func Context(t *testing.T, dur time.Duration) context.Context {
 	t.Cleanup(cancel)
 	return ctx
 }
-
-func RequireRecvCtx[A any](ctx context.Context, t testing.TB, c <-chan A) (a A) {
-	t.Helper()
-	select {
-	case <-ctx.Done():
-		t.Fatal("timeout")
-		return a
-	case a = <-c:
-		return a
-	}
-}
-
-// NOTE: no AssertRecvCtx because it'd be bad if we returned a default value on
-// the cases it times out.
-
-func RequireSendCtx[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
-	t.Helper()
-	select {
-	case <-ctx.Done():
-		t.Fatal("timeout")
-	case c <- a:
-		// OK!
-	}
-}
-
-func AssertSendCtx[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
-	t.Helper()
-	select {
-	case <-ctx.Done():
-		t.Error("timeout")
-	case c <- a:
-		// OK!
-	}
-}