@@ -21,13 +21,17 @@ import (
2121 "errors"
2222 "log/slog"
2323 "strings"
24+ "sync"
2425 "testing"
2526 "time"
2627
2728 "github.com/agent-substrate/substrate/pkg/proto/ateapipb"
2829 corev3 "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
2930 extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
3031 envoy_type "github.com/envoyproxy/go-control-plane/envoy/type/v3"
32+ "go.opentelemetry.io/otel"
33+ "go.opentelemetry.io/otel/propagation"
34+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
3135 "google.golang.org/grpc"
3236 "google.golang.org/grpc/codes"
3337 "google.golang.org/grpc/status"
@@ -230,17 +234,19 @@ func TestExtProcHeadersEvaluation(t *testing.T) {
230234 }
231235
232236 mutation := res .Response .GetHeaderMutation ()
233- if len (mutation .GetSetHeaders ()) != 1 {
234- t .Fatalf ("expected exactly one Header option set, found: %v" , mutation .GetSetHeaders ())
235- }
236-
237- headerOption := mutation .GetSetHeaders ()[0 ]
238- if strings .ToLower (headerOption .Header .Key ) != ":authority" {
239- t .Errorf ("invalid resulting dynamic parameter key: %s" , headerOption .Header .Key )
237+ headers := mutation .GetSetHeaders ()
238+ var found bool
239+ for _ , h := range headers {
240+ if strings .ToLower (h .Header .Key ) == ":authority" {
241+ found = true
242+ if string (h .Header .RawValue ) != tc .expectedTarget {
243+ t .Errorf ("invalid destination mapping found: %s, expected: %s" , h .Header .RawValue , tc .expectedTarget )
244+ }
245+ break
246+ }
240247 }
241-
242- if string (headerOption .Header .RawValue ) != tc .expectedTarget {
243- t .Errorf ("invalid destination mapping found: %s, expected: %s" , headerOption .Header .RawValue , tc .expectedTarget )
248+ if ! found {
249+ t .Errorf (":authority header not found in mutation, got: %v" , headers )
244250 }
245251
246252 // Confirm that query logs recorded metric trace details
@@ -252,3 +258,100 @@ func TestExtProcHeadersEvaluation(t *testing.T) {
252258 })
253259 }
254260}
261+
262+ // inMemoryExporter stores ended spans for test inspection.
263+ type inMemoryExporter struct {
264+ mu sync.Mutex
265+ spans []sdktrace.ReadOnlySpan
266+ }
267+
268+ func (e * inMemoryExporter ) ExportSpans (_ context.Context , spans []sdktrace.ReadOnlySpan ) error {
269+ e .mu .Lock ()
270+ defer e .mu .Unlock ()
271+ e .spans = append (e .spans , spans ... )
272+ return nil
273+ }
274+
275+ func (e * inMemoryExporter ) Shutdown (context.Context ) error {
276+ e .mu .Lock ()
277+ defer e .mu .Unlock ()
278+ e .spans = nil
279+ return nil
280+ }
281+
282+ func (e * inMemoryExporter ) Ended () []sdktrace.ReadOnlySpan {
283+ e .mu .Lock ()
284+ defer e .mu .Unlock ()
285+ ret := make ([]sdktrace.ReadOnlySpan , len (e .spans ))
286+ copy (ret , e .spans )
287+ return ret
288+ }
289+
290+ func TestInjectTraceContext_AddsTraceparentMatchingParentSpan (t * testing.T ) {
291+ prevProp := otel .GetTextMapPropagator ()
292+ otel .SetTextMapPropagator (propagation.TraceContext {})
293+ t .Cleanup (func () { otel .SetTextMapPropagator (prevProp ) })
294+
295+ exporter := & inMemoryExporter {}
296+ tp := sdktrace .NewTracerProvider (
297+ sdktrace .WithSampler (sdktrace .AlwaysSample ()),
298+ sdktrace .WithSpanProcessor (sdktrace .NewSimpleSpanProcessor (exporter )),
299+ )
300+
301+ ctx , span := tp .Tracer ("test" ).Start (context .Background (), "parent" )
302+ parentTraceID := span .SpanContext ().TraceID ()
303+
304+ mutation := & extprocv3.HeaderMutation {}
305+ injectTraceContext (ctx , mutation )
306+ span .End ()
307+
308+ var traceparent string
309+ for _ , h := range mutation .GetSetHeaders () {
310+ if strings .ToLower (h .Header .Key ) == "traceparent" {
311+ traceparent = h .Header .Value
312+ if h .GetAppendAction () != corev3 .HeaderValueOption_OVERWRITE_IF_EXISTS_OR_ADD {
313+ t .Errorf ("AppendAction = %v, want OVERWRITE_IF_EXISTS_OR_ADD" , h .GetAppendAction ())
314+ }
315+ break
316+ }
317+ }
318+ if traceparent == "" {
319+ t .Fatal ("traceparent not found in header mutation" )
320+ }
321+
322+ parts := strings .Split (traceparent , "-" )
323+ if len (parts ) < 2 {
324+ t .Fatalf ("invalid traceparent format: %s" , traceparent )
325+ }
326+ if parts [1 ] != parentTraceID .String () {
327+ t .Errorf ("trace ID in traceparent = %s, want parent trace ID = %s" , parts [1 ], parentTraceID .String ())
328+ }
329+ }
330+
331+ func TestInjectTraceContext_AppendActionDefaultsToOverwrite (t * testing.T ) {
332+ prevProp := otel .GetTextMapPropagator ()
333+ otel .SetTextMapPropagator (propagation.TraceContext {})
334+ t .Cleanup (func () { otel .SetTextMapPropagator (prevProp ) })
335+
336+ exporter := & inMemoryExporter {}
337+ tp := sdktrace .NewTracerProvider (
338+ sdktrace .WithSampler (sdktrace .AlwaysSample ()),
339+ sdktrace .WithSpanProcessor (sdktrace .NewSimpleSpanProcessor (exporter )),
340+ )
341+
342+ ctx , span := tp .Tracer ("test" ).Start (context .Background (), "parent" )
343+ defer span .End ()
344+
345+ mutation := & extprocv3.HeaderMutation {}
346+ injectTraceContext (ctx , mutation )
347+
348+ for _ , h := range mutation .GetSetHeaders () {
349+ if strings .ToLower (h .Header .Key ) == "traceparent" {
350+ if h .GetAppendAction () != corev3 .HeaderValueOption_OVERWRITE_IF_EXISTS_OR_ADD {
351+ t .Errorf ("AppendAction = %v, want OVERWRITE_IF_EXISTS_OR_ADD" , h .GetAppendAction ())
352+ }
353+ return
354+ }
355+ }
356+ t .Error ("traceparent not found in mutation" )
357+ }
0 commit comments