@@ -32,6 +32,7 @@ import (
3232 "google.golang.org/grpc/metadata"
3333 "google.golang.org/grpc/status"
3434
35+ "github.com/open-telemetry/opentelemetry-collector-contrib/internal/grpcutil"
3536 "github.com/open-telemetry/opentelemetry-collector-contrib/internal/otelarrow/netstats"
3637 "github.com/open-telemetry/opentelemetry-collector-contrib/internal/otelarrow/testdata"
3738)
@@ -576,65 +577,94 @@ func TestArrowExporterStreaming(t *testing.T) {
576577
577578// TestArrowExporterHeaders tests a mix of outgoing context headers.
578579func TestArrowExporterHeaders (t * testing.T ) {
579- tc := newSingleStreamMetadataTestCase ( t )
580- channel := newHealthyTestChannel ()
580+ for _ , withDeadline := range [] bool { true , false } {
581+ t . Run ( fmt . Sprint ( "with_deadline=" , withDeadline ), func ( t * testing. T ) {
581582
582- tc .traceCall .AnyTimes ().DoAndReturn (tc .returnNewStream (channel ))
583+ tc := newSingleStreamMetadataTestCase (t )
584+ channel := newHealthyTestChannel ()
583585
584- ctx , cancel := context .WithCancel (context .Background ())
585- defer cancel ()
586- require .NoError (t , tc .exporter .Start (ctx ))
586+ tc .traceCall .AnyTimes ().DoAndReturn (tc .returnNewStream (channel ))
587587
588- var expectOutput []metadata. MD
589- var actualOutput []metadata. MD
588+ ctx , cancel := context . WithCancel ( context . Background ())
589+ defer cancel ()
590590
591- var wg sync.WaitGroup
592- wg .Add (1 )
593- go func () {
594- defer wg .Done ()
595- md := metadata.MD {}
596- hpd := hpack .NewDecoder (4096 , func (f hpack.HeaderField ) {
597- md [f .Name ] = append (md [f .Name ], f .Value )
598- })
599- for data := range channel .sendChannel () {
600- if len (data .Headers ) == 0 {
601- actualOutput = append (actualOutput , nil )
602- } else {
603- _ , err := hpd .Write (data .Headers )
591+ require .NoError (t , tc .exporter .Start (ctx ))
592+
593+ var expectOutput []metadata.MD
594+ var actualOutput []metadata.MD
595+
596+ var wg sync.WaitGroup
597+ wg .Add (1 )
598+ go func () {
599+ defer wg .Done ()
600+ md := metadata.MD {}
601+ hpd := hpack .NewDecoder (4096 , func (f hpack.HeaderField ) {
602+ md [f .Name ] = append (md [f .Name ], f .Value )
603+ })
604+ for data := range channel .sendChannel () {
605+ if len (data .Headers ) == 0 {
606+ actualOutput = append (actualOutput , nil )
607+ } else {
608+ _ , err := hpd .Write (data .Headers )
609+ require .NoError (t , err )
610+ actualOutput = append (actualOutput , md )
611+ md = metadata.MD {}
612+ }
613+ channel .recv <- statusOKFor (data .BatchId )
614+ }
615+ }()
616+
617+ for times := 0 ; times < 10 ; times ++ {
618+ input := testdata .GenerateTraces (2 )
619+
620+ if times % 2 == 1 {
621+ md := metadata.MD {
622+ "expected1" : []string {"metadata1" },
623+ "expected2" : []string {fmt .Sprint (times )},
624+ "otlp-pdata-size" : []string {"329" },
625+ }
626+ expectOutput = append (expectOutput , md )
627+ } else {
628+ expectOutput = append (expectOutput , metadata.MD {
629+ "otlp-pdata-size" : []string {"329" },
630+ })
631+ }
632+
633+ sendCtx := ctx
634+ if withDeadline {
635+ var sendCancel context.CancelFunc
636+ sendCtx , sendCancel = context .WithTimeout (sendCtx , time .Second )
637+ defer sendCancel ()
638+ }
639+
640+ sent , err := tc .exporter .SendAndWait (sendCtx , input )
604641 require .NoError (t , err )
605- actualOutput = append (actualOutput , md )
606- md = metadata.MD {}
642+ require .True (t , sent )
607643 }
608- channel .recv <- statusOKFor (data .BatchId )
609- }
610- }()
611-
612- for times := 0 ; times < 10 ; times ++ {
613- input := testdata .GenerateTraces (2 )
644+ // Stop the test conduit started above.
645+ cancel ()
646+ wg .Wait ()
614647
615- if times % 2 == 1 {
616- md := metadata.MD {
617- "expected1" : []string {"metadata1" },
618- "expected2" : []string {fmt .Sprint (times )},
619- "otlp-pdata-size" : []string {"329" },
648+ // Manual check for proper deadline propagation. Since the test
649+ // is timed we don't expect an exact match.
650+ if withDeadline {
651+ for _ , out := range actualOutput {
652+ dead := out .Get ("grpc-timeout" )
653+ require .Len (t , dead , 1 )
654+ require .NotEmpty (t , dead [0 ])
655+ to , err := grpcutil .DecodeTimeout (dead [0 ])
656+ require .NoError (t , err )
657+ // Allow the test to lapse for 0.5s.
658+ require .Less (t , time .Second / 2 , to )
659+ require .GreaterOrEqual (t , time .Second , to )
660+ out .Delete ("grpc-timeout" )
661+ }
620662 }
621- expectOutput = append (expectOutput , md )
622- } else {
623- expectOutput = append (expectOutput , metadata.MD {
624- "otlp-pdata-size" : []string {"329" },
625- })
626- }
627663
628- sent , err := tc . exporter . SendAndWait ( context . Background (), input )
629- require .NoError (t , err )
630- require . True ( t , sent )
664+ require . Equal ( t , expectOutput , actualOutput )
665+ require .NoError (t , tc . exporter . Shutdown ( ctx ) )
666+ } )
631667 }
632- // Stop the test conduit started above.
633- cancel ()
634- wg .Wait ()
635-
636- require .Equal (t , expectOutput , actualOutput )
637- require .NoError (t , tc .exporter .Shutdown (ctx ))
638668}
639669
640670// TestArrowExporterIsTraced tests whether trace and span ID are
0 commit comments