use helper functions

hugodutka · hugodutka · commit 39cbc5b63b7e · 2025-05-22T11:04:20.000Z
diff --git a/coderd/telemetry/telemetry_test.go b/coderd/telemetry/telemetry_test.go
@@ -450,10 +450,10 @@ func mockTelemetryServer(ctx context.Context, t *testing.T) (*url.URL, chan *tel
 		dd := &telemetry.Deployment{}
 		err := json.NewDecoder(r.Body).Decode(dd)
 		require.NoError(t, err)
-		select {
-		case <-ctx.Done():
-			t.Fatal("timed out sending deployment")
-		case deployment <- dd:
+		ok := testutil.AssertSend(ctx, t, deployment, dd)
+		if !ok {
+			w.WriteHeader(http.StatusInternalServerError)
+			return
 		}
 		// Ensure the header is sent only after deployment is sent
 		w.WriteHeader(http.StatusAccepted)
@@ -463,10 +463,10 @@ func mockTelemetryServer(ctx context.Context, t *testing.T) (*url.URL, chan *tel
 		ss := &telemetry.Snapshot{}
 		err := json.NewDecoder(r.Body).Decode(ss)
 		require.NoError(t, err)
-		select {
-		case <-ctx.Done():
-			t.Fatal("timed out sending snapshot")
-		case snapshot <- ss:
+		ok := testutil.AssertSend(ctx, t, snapshot, ss)
+		if !ok {
+			w.WriteHeader(http.StatusInternalServerError)
+			return
 		}
 		// Ensure the header is sent only after snapshot is sent
 		w.WriteHeader(http.StatusAccepted)
@@ -487,7 +487,7 @@ func collectSnapshot(
 ) (*telemetry.Deployment, *telemetry.Snapshot) {
 	t.Helper()
 
-	serverURL, deploymentChan, snapshotChan := mockTelemetryServer(ctx, t)
+	serverURL, deployment, snapshot := mockTelemetryServer(ctx, t)
 
 	options := telemetry.Options{
 		Database:     db,
@@ -503,19 +503,5 @@ func collectSnapshot(
 	require.NoError(t, err)
 	t.Cleanup(reporter.Close)
 
-	var deployment *telemetry.Deployment
-	var snapshot *telemetry.Snapshot
-
-	select {
-	case <-ctx.Done():
-		t.Fatal("timed out collecting deployment")
-	case deployment = <-deploymentChan:
-	}
-	select {
-	case <-ctx.Done():
-		t.Fatal("timed out collecting snapshot")
-	case snapshot = <-snapshotChan:
-	}
-
-	return deployment, snapshot
+	return testutil.RequireReceive(ctx, t, deployment), testutil.RequireReceive(ctx, t, snapshot)
 }
diff --git a/testutil/chan.go b/testutil/chan.go
@@ -55,3 +55,61 @@ func RequireSend[A any](ctx context.Context, t testing.TB, c chan<- A, a A) {
 		// OK!
 	}
 }
+
+// SoftTryReceive will attempt to receive a value from the chan and return it. If
+// the context expires before a value can be received, it will mark the test as
+// failed but continue execution. If the channel is closed, the zero value of the
+// channel type will be returned.
+// The second return value indicates whether the receive was successful. In
+// particular, if the channel is closed, the second return value will be true.
+//
+// Safety: can be called from any goroutine.
+func SoftTryReceive[A any](ctx context.Context, t testing.TB, c <-chan A) (A, bool) {
+	t.Helper()
+	select {
+	case <-ctx.Done():
+		t.Error("timeout")
+		var a A
+		return a, false
+	case a := <-c:
+		return a, true
+	}
+}
+
+// AssertReceive 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 mark the test as failed but continue execution.
+// The second return value indicates whether the receive was successful.
+//
+// Safety: can be called from any goroutine.
+func AssertReceive[A any](ctx context.Context, t testing.TB, c <-chan A) (A, bool) {
+	t.Helper()
+	select {
+	case <-ctx.Done():
+		t.Error("timeout")
+		var a A
+		return a, false
+	case a, ok := <-c:
+		if !ok {
+			t.Error("channel closed")
+		}
+		return a, ok
+	}
+}
+
+// AssertSend will send the given value over the chan and then return. If
+// the context expires before the send succeeds, it will mark the test as failed
+// but continue execution.
+// The second return value indicates whether the send was successful.
+//
+// Safety: can be called from any goroutine.
+func AssertSend[A any](ctx context.Context, t testing.TB, c chan<- A, a A) bool {
+	t.Helper()
+	select {
+	case <-ctx.Done():
+		t.Error("timeout")
+		return false
+	case c <- a:
+		return true
+	}
+}
