Add helpers for testing exported spans (#11655)

Introduce reusable test helpers for dealing with OTEL spans.

common/testing/testtelemetry/spans.go

@@ -0,0 +1,104 @@
+package testtelemetry
+
+import (
+	"slices"
+
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+	oteltrace "go.opentelemetry.io/otel/trace"
+)
+
+// LocalSpanID contains stable identifiers for one span and its trace-local relationships.
+type LocalSpanID struct {
+	Trace  int
+	Span   int
+	Parent int
+}
+
+// idGenerator assigns stable, one-based local IDs while preserving zero as no ID.
+type idGenerator[T comparable] struct {
+	ids map[T]int
+}
+
+func (g *idGenerator[T]) get(value T) int {
+	var zero T
+	if value == zero {
+		return 0
+	}
+	if id, ok := g.ids[value]; ok {
+		return id
+	}
+	if g.ids == nil {
+		g.ids = make(map[T]int)
+	}
+	id := len(g.ids) + 1
+	g.ids[value] = id
+	return id
+}
+
+func (g *idGenerator[T]) lookup(value T) int {
+	return g.ids[value]
+}
+
+func localIDs[T comparable](values []T) []int {
+	var ids idGenerator[T]
+	result := make([]int, len(values))
+	for i, value := range values {
+		result[i] = ids.get(value)
+	}
+	return result
+}
+
+// LocalAttributeIDs assigns stable IDs to an attribute's values in span order.
+func LocalAttributeIDs(spans tracetest.SpanStubs, key attribute.Key) []int {
+	values := make([]string, len(spans))
+	for i, span := range spans {
+		if value, ok := SpanAttribute(span, key); ok {
+			values[i] = value.String()
+		}
+	}
+	return localIDs(values)
+}
+
+// LocalSpanIDs assigns stable, one-based IDs in the order spans are provided.
+// Callers must sort spans first when they need deterministic ordering.
+func LocalSpanIDs(spans tracetest.SpanStubs) []LocalSpanID {
+	var traceIDs idGenerator[oteltrace.TraceID]
+	spanIDs := make(map[oteltrace.TraceID]*idGenerator[oteltrace.SpanID])
+	for _, span := range spans {
+		traceID := span.SpanContext.TraceID()
+		traceIDs.get(traceID)
+		if spanIDs[traceID] == nil {
+			spanIDs[traceID] = new(idGenerator[oteltrace.SpanID])
+		}
+		spanIDs[traceID].get(span.SpanContext.SpanID())
+	}
+
+	result := make([]LocalSpanID, len(spans))
+	for i, span := range spans {
+		traceID := span.SpanContext.TraceID()
+		result[i] = LocalSpanID{
+			Trace:  traceIDs.lookup(traceID),
+			Span:   spanIDs[traceID].lookup(span.SpanContext.SpanID()),
+			Parent: spanIDs[traceID].lookup(span.Parent.SpanID()),
+		}
+	}
+	return result
+}
+
+func SpanAttribute(span tracetest.SpanStub, key attribute.Key) (attribute.Value, bool) {
+	for _, attr := range span.Attributes {
+		if attr.Key == key {
+			return attr.Value, true
+		}
+	}
+	return attribute.Value{}, false
+}
+
+// FilterSpans returns matching spans without modifying the input slice.
+func FilterSpans(spans tracetest.SpanStubs, keep func(tracetest.SpanStub) bool) tracetest.SpanStubs {
+	filtered := slices.Clone(spans)
+	return slices.DeleteFunc(filtered, func(span tracetest.SpanStub) bool {
+		return !keep(span)
+	})
+}

common/testing/testtelemetry/spans_test.go

@@ -0,0 +1,67 @@
+package testtelemetry
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+	oteltrace "go.opentelemetry.io/otel/trace"
+)
+
+func TestLocalAttributeIDs(t *testing.T) {
+	t.Parallel()
+
+	spans := tracetest.SpanStubs{
+		{},
+		{Attributes: []attribute.KeyValue{attribute.String("key", "first")}},
+		{Attributes: []attribute.KeyValue{attribute.String("key", "second")}},
+		{Attributes: []attribute.KeyValue{attribute.String("key", "first")}},
+	}
+	require.Equal(t, []int{0, 1, 2, 1}, LocalAttributeIDs(spans, "key"))
+}
+
+func TestLocalSpanIDs(t *testing.T) {
+	t.Parallel()
+
+	traceID1 := oteltrace.TraceID{1}
+	traceID2 := oteltrace.TraceID{2}
+	parentSpanID := oteltrace.SpanID{1}
+	childSpanID := oteltrace.SpanID{2}
+	spanContext := func(traceID oteltrace.TraceID, spanID oteltrace.SpanID) oteltrace.SpanContext {
+		return oteltrace.NewSpanContext(oteltrace.SpanContextConfig{TraceID: traceID, SpanID: spanID})
+	}
+	spans := tracetest.SpanStubs{
+		{SpanContext: spanContext(traceID1, childSpanID), Parent: spanContext(traceID1, parentSpanID)},
+		{SpanContext: spanContext(traceID1, parentSpanID)},
+		{SpanContext: spanContext(traceID2, parentSpanID)},
+	}
+	require.Equal(t, []LocalSpanID{
+		{Trace: 1, Span: 1, Parent: 2},
+		{Trace: 1, Span: 2},
+		{Trace: 2, Span: 1},
+	}, LocalSpanIDs(spans))
+}
+
+func TestSpanAttribute(t *testing.T) {
+	t.Parallel()
+
+	span := tracetest.SpanStub{Attributes: []attribute.KeyValue{attribute.String("key", "value")}}
+	value, ok := SpanAttribute(span, "key")
+	require.True(t, ok)
+	require.Equal(t, "value", value.AsString())
+
+	_, ok = SpanAttribute(span, "missing")
+	require.False(t, ok)
+}
+
+func TestFilterSpans(t *testing.T) {
+	t.Parallel()
+
+	spans := tracetest.SpanStubs{{Name: "keep"}, {Name: "drop"}}
+	filtered := FilterSpans(spans, func(span tracetest.SpanStub) bool {
+		return span.Name == "keep"
+	})
+	require.Equal(t, tracetest.SpanStubs{{Name: "keep"}}, filtered)
+	require.Len(t, spans, 2)
+}

tests/nexus_otel_test.go

@@ -30,6 +30,7 @@ import (
 	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/nexus/nexusrpc"
 	"go.temporal.io/server/common/testing/parallelsuite"
+	"go.temporal.io/server/common/testing/testtelemetry"
 	"go.temporal.io/server/tests/testcore"
 	"google.golang.org/protobuf/types/known/durationpb"
 )
@@ -396,39 +397,27 @@ func (s *NexusOTELSuite) requireNexusHTTPSpans(
 ) tracetest.SpanStubs {
 	s.T().Helper()
 	var httpSpans tracetest.SpanStubs
-	s.Await(func(s *NexusOTELSuite) {
+	requireExportedSpans(s, exporter, expected, func(spans tracetest.SpanStubs) []nexusHTTPSpan {
 		var actual []nexusHTTPSpan
-		actual, httpSpans = s.nexusHTTPSpans(exporter.GetSpans())
-		s.Require().Len(actual, len(expected))
-		s.Require().Equal(expected, actual)
-	}, 10*time.Second, 100*time.Millisecond)
+		actual, httpSpans = s.nexusHTTPSpans(spans)
+		return actual
+	})
 	return httpSpans
 }
 
 func (s *NexusOTELSuite) nexusHTTPSpans(
 	spans tracetest.SpanStubs,
 ) ([]nexusHTTPSpan, tracetest.SpanStubs) {
-	httpSpans := slices.DeleteFunc(spans, func(span tracetest.SpanStub) bool {
-		return span.InstrumentationScope.Name != otelhttp.ScopeName
+	httpSpans := testtelemetry.FilterSpans(spans, func(span tracetest.SpanStub) bool {
+		return span.InstrumentationScope.Name == otelhttp.ScopeName
 	})
 	slices.SortFunc(httpSpans, func(a, b tracetest.SpanStub) int {
 		return a.StartTime.Compare(b.StartTime)
 	})
-
-	traceIDs := make(map[oteltrace.TraceID]int)
-	spanIDs := make(map[oteltrace.TraceID]map[oteltrace.SpanID]int)
-	for _, span := range httpSpans {
-		traceID := span.SpanContext.TraceID()
-		if _, ok := traceIDs[traceID]; !ok {
-			traceIDs[traceID] = len(traceIDs) + 1
-			spanIDs[traceID] = make(map[oteltrace.SpanID]int)
-		}
-		spanIDs[traceID][span.SpanContext.SpanID()] = len(spanIDs[traceID]) + 1
-	}
+	localIDs := testtelemetry.LocalSpanIDs(httpSpans)
 
 	result := make([]nexusHTTPSpan, 0, len(httpSpans))
-	for _, span := range httpSpans {
-		traceID := span.SpanContext.TraceID()
+	for i, span := range httpSpans {
 		var serviceName string
 		if span.Resource != nil {
 			if value, ok := span.Resource.Set().Value(semconv.ServiceNameKey); ok {
@@ -455,9 +444,9 @@ func (s *NexusOTELSuite) nexusHTTPSpans(
 			}
 		}
 		result = append(result, nexusHTTPSpan{
-			TraceID:      traceIDs[traceID],
-			SpanID:       spanIDs[traceID][span.SpanContext.SpanID()],
-			ParentSpanID: spanIDs[traceID][span.Parent.SpanID()],
+			TraceID:      localIDs[i].Trace,
+			SpanID:       localIDs[i].Span,
+			ParentSpanID: localIDs[i].Parent,
 			Name:         span.Name,
 			ServiceName:  serviceName,
 			Kind:         span.SpanKind,
@@ -468,3 +457,17 @@ func (s *NexusOTELSuite) nexusHTTPSpans(
 	}
 	return result, httpSpans
 }
+
+func requireExportedSpans[T any](
+	s *NexusOTELSuite,
+	exporter *tracetest.InMemoryExporter,
+	expected []T,
+	project func(tracetest.SpanStubs) []T,
+) {
+	s.T().Helper()
+	s.Await(func(s *NexusOTELSuite) {
+		actual := project(exporter.GetSpans())
+		s.Require().Len(actual, len(expected))
+		s.Require().Equal(expected, actual)
+	}, 10*time.Second, 100*time.Millisecond)
+}