github.com/blend/go-sdk@v1.20220411.3/tracing/context_test.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package tracing
     9  
    10  import (
    11  	"context"
    12  	"testing"
    13  
    14  	"github.com/blend/go-sdk/assert"
    15  	"github.com/blend/go-sdk/logger"
    16  
    17  	opentracing "github.com/opentracing/opentracing-go"
    18  )
    19  
    20  type spanContextWithoutGetters struct {
    21  	opentracing.SpanContext
    22  }
    23  
    24  type spanContextWithSpanID struct {
    25  	opentracing.SpanContext
    26  	spanID uint64
    27  }
    28  
    29  func (c spanContextWithSpanID) SpanID() uint64 {
    30  	return c.spanID
    31  }
    32  
    33  type spanContextWithTraceID struct {
    34  	opentracing.SpanContext
    35  	traceID uint64
    36  }
    37  
    38  func (c spanContextWithTraceID) TraceID() uint64 {
    39  	return c.traceID
    40  }
    41  
    42  type spanContextWithAllGetters struct {
    43  	opentracing.SpanContext
    44  	spanID  uint64
    45  	traceID uint64
    46  }
    47  
    48  func (c spanContextWithAllGetters) SpanID() uint64 {
    49  	return c.spanID
    50  }
    51  
    52  func (c spanContextWithAllGetters) TraceID() uint64 {
    53  	return c.traceID
    54  }
    55  
    56  func TestWithTraceAnnotations_NoGetters(t *testing.T) {
    57  	assert := assert.New(t)
    58  
    59  	ctx := context.Background()
    60  	ctx = WithTraceAnnotations(ctx, spanContextWithoutGetters{})
    61  	assert.Empty(logger.GetAnnotations(ctx))
    62  }
    63  
    64  func TestWithTraceAnnotations_AllGetters(t *testing.T) {
    65  	assert := assert.New(t)
    66  
    67  	ctx := context.Background()
    68  	ctx = WithTraceAnnotations(ctx, spanContextWithAllGetters{
    69  		spanID:  123,
    70  		traceID: 456,
    71  	})
    72  
    73  	annotations := logger.GetAnnotations(ctx)
    74  	assert.Len(annotations, 2)
    75  	assert.Equal("123", annotations[LoggerAnnotationTracingSpanID])
    76  	assert.Equal("456", annotations[LoggerAnnotationTracingTraceID])
    77  }
    78  
    79  func TestWithTraceAnnotations_SpanIDProvider(t *testing.T) {
    80  	assert := assert.New(t)
    81  
    82  	ctx := context.Background()
    83  	ctx = WithTraceAnnotations(ctx, spanContextWithSpanID{
    84  		spanID: 123,
    85  	})
    86  
    87  	annotations := logger.GetAnnotations(ctx)
    88  	assert.Len(annotations, 1)
    89  	assert.Equal("123", annotations[LoggerAnnotationTracingSpanID])
    90  }
    91  
    92  func TestWithTraceAnnotations_TraceIDProvider(t *testing.T) {
    93  	assert := assert.New(t)
    94  
    95  	ctx := context.Background()
    96  	ctx = WithTraceAnnotations(ctx, spanContextWithTraceID{
    97  		traceID: 456,
    98  	})
    99  
   100  	annotations := logger.GetAnnotations(ctx)
   101  	assert.Len(annotations, 1)
   102  	assert.Equal("456", annotations[LoggerAnnotationTracingTraceID])
   103  }