@@ -29,7 +29,7 @@ import (
2929 otelpropagators "go.opentelemetry.io/otel/propagators"
3030 "go.opentelemetry.io/otel/semconv"
3131
32- mocktracer "go.opentelemetry.io/contrib/internal /trace"
32+ "go.opentelemetry.io/otel/api /trace/tracetest "
3333)
3434
3535const (
@@ -39,7 +39,8 @@ const (
3939func TestWrapPartitionConsumer (t * testing.T ) {
4040 propagators := otelpropagators.TraceContext {}
4141 // Mock provider
42- provider , mt := NewTracerProviderAndTracer ()
42+ sr := new (tracetest.StandardSpanRecorder )
43+ provider := tracetest .NewTracerProvider (tracetest .WithSpanRecorder (sr ))
4344
4445 // Mock partition consumer controller
4546 consumer := mocks .NewConsumer (t , sarama .NewConfig ())
@@ -51,13 +52,14 @@ func TestWrapPartitionConsumer(t *testing.T) {
5152
5253 partitionConsumer = WrapPartitionConsumer (partitionConsumer , WithTracerProvider (provider ), WithPropagators (propagators ))
5354
54- consumeAndCheck (t , mt , mockPartitionConsumer , partitionConsumer )
55+ consumeAndCheck (t , provider . Tracer ( defaultTracerName ), sr . Completed , mockPartitionConsumer , partitionConsumer )
5556}
5657
5758func TestWrapConsumer (t * testing.T ) {
5859 propagators := otelpropagators.TraceContext {}
5960 // Mock provider
60- provider , mt := NewTracerProviderAndTracer ()
61+ sr := new (tracetest.StandardSpanRecorder )
62+ provider := tracetest .NewTracerProvider (tracetest .WithSpanRecorder (sr ))
6163
6264 // Mock partition consumer controller
6365 mockConsumer := mocks .NewConsumer (t , sarama .NewConfig ())
@@ -70,10 +72,10 @@ func TestWrapConsumer(t *testing.T) {
7072 partitionConsumer , err := consumer .ConsumePartition (topic , 0 , 0 )
7173 require .NoError (t , err )
7274
73- consumeAndCheck (t , mt , mockPartitionConsumer , partitionConsumer )
75+ consumeAndCheck (t , provider . Tracer ( defaultTracerName ), sr . Completed , mockPartitionConsumer , partitionConsumer )
7476}
7577
76- func consumeAndCheck (t * testing.T , mt * mocktracer .Tracer , mockPartitionConsumer * mocks.PartitionConsumer , partitionConsumer sarama.PartitionConsumer ) {
78+ func consumeAndCheck (t * testing.T , mt trace .Tracer , complFn func () [] * tracetest. Span , mockPartitionConsumer * mocks.PartitionConsumer , partitionConsumer sarama.PartitionConsumer ) {
7779 // Create message with span context
7880 ctx , _ := mt .Start (context .Background (), "" )
7981 message := sarama.ConsumerMessage {Key : []byte ("foo" )}
@@ -93,7 +95,7 @@ func consumeAndCheck(t *testing.T, mt *mocktracer.Tracer, mockPartitionConsumer
9395 <- partitionConsumer .Messages ()
9496
9597 // Check spans length
96- spans := mt . EndedSpans ()
98+ spans := complFn ()
9799 assert .Len (t , spans , 2 )
98100
99101 expectedList := []struct {
@@ -133,17 +135,17 @@ func consumeAndCheck(t *testing.T, mt *mocktracer.Tracer, mockPartitionConsumer
133135 t .Run (fmt .Sprint ("index" , i ), func (t * testing.T ) {
134136 span := spans [i ]
135137
136- assert .Equal (t , expected .parentSpanID , span .ParentSpanID )
138+ assert .Equal (t , expected .parentSpanID , span .ParentSpanID () )
137139
138140 remoteSpanFromMessage := trace .RemoteSpanContextFromContext (propagators .Extract (context .Background (), NewConsumerMessageCarrier (msgList [i ])))
139141 assert .Equal (t , span .SpanContext (), remoteSpanFromMessage ,
140142 "span context should be injected into the consumer message headers" )
141143
142- assert .Equal (t , "kafka.consume" , span .Name )
143- assert .Equal (t , expected .kind , span .Kind )
144+ assert .Equal (t , "kafka.consume" , span .Name () )
145+ assert .Equal (t , expected .kind , span .SpanKind () )
144146 assert .Equal (t , expected .msgKey , msgList [i ].Key )
145147 for _ , k := range expected .labelList {
146- assert .Equal (t , k .Value , span .Attributes [k .Key ], k .Key )
148+ assert .Equal (t , k .Value , span .Attributes () [k .Key ], k .Key )
147149 }
148150 })
149151 }
@@ -164,7 +166,7 @@ func TestConsumerConsumePartitionWithError(t *testing.T) {
164166
165167func BenchmarkWrapPartitionConsumer (b * testing.B ) {
166168 // Mock provider
167- provider , _ := NewTracerProviderAndTracer ()
169+ provider := tracetest . NewTracerProvider ()
168170
169171 mockPartitionConsumer , partitionConsumer := createMockPartitionConsumer (b )
170172
0 commit comments