Add a collectStreamedTrace test helper

Spans are buffered per trace before they flush, so an earlier page load's spans
can still be arriving when a test starts collecting. This resolves with only the
spans sharing the matched segment's trace, so tests asserting on one request's
children cannot pick up leftovers.

Author: andreiborza

dev-packages/test-utils/src/event-proxy-server.ts

@@ -637,6 +637,58 @@ export function collectStreamedSpans(
   }).then(() => collected);
 }
 
+/**
+ * Accumulate streamed spans until a segment span matches `isSegment`, then resolve with only the
+ * spans that share that segment's trace.
+ *
+ * Spans are buffered per trace before they flush, so spans from an earlier page load can still be
+ * arriving when a test starts collecting. Scoping to a single trace is what keeps those leftovers
+ * from ending the wait early or from failing an assertion about the trace under test. Prefer this
+ * over {@link collectStreamedSpans} whenever a test asserts on the children of one specific request.
+ *
+ * `isDone` receives the spans collected so far for that trace and defaults to "the segment span has
+ * arrived". Pass it when a test needs particular children, since a child can flush before or after
+ * its segment.
+ *
+ * @example
+ * ```ts
+ * const spans = await collectStreamedTrace(PROXY_SERVER_NAME, span => span.name === 'GET /redis');
+ * expect(spans.map(span => span.name)).toContainEqual('redis-set');
+ * ```
+ *
+ * @example
+ * ```ts
+ * // Wait for the segment and at least two of its db spans
+ * const spans = await collectStreamedTrace(
+ *   PROXY_SERVER_NAME,
+ *   span => span.name === 'GET /redis',
+ *   spansOfTrace => spansOfTrace.filter(span => getSpanOp(span) === 'db.query').length >= 2,
+ * );
+ * ```
+ */
+export function collectStreamedTrace(
+  proxyServerName: string,
+  isSegment: (span: SerializedStreamedSpan) => boolean,
+  isDone?: (spansOfTrace: SerializedStreamedSpan[]) => boolean,
+): Promise<SerializedStreamedSpan[]> {
+  const collected: SerializedStreamedSpan[] = [];
+  let traceId: string | undefined;
+
+  return waitForStreamedSpans(proxyServerName, spans => {
+    collected.push(...spans);
+
+    // Only the first matching segment counts, so a later page load on the same route cannot move
+    // the collection to a different trace halfway through.
+    traceId ??= collected.find(span => span.is_segment && isSegment(span))?.trace_id;
+
+    if (traceId === undefined) {
+      return false;
+    }
+
+    return isDone ? isDone(collected.filter(span => span.trace_id === traceId)) : true;
+  }).then(() => collected.filter(span => span.trace_id === traceId));
+}
+
 /**
  * Helper to get the span operation from a Span V2 JSON object.
  *

dev-packages/test-utils/src/index.ts

@@ -12,6 +12,7 @@ export {
   waitForStreamedSpans,
   waitForStreamedSpanEnvelope,
   collectStreamedSpans,
+  collectStreamedTrace,
   getSpanOp,
 } from './event-proxy-server';
 export type { SerializedStreamedSpan } from '@sentry/core';