mocktracer.go

Documentation: github.com/opentracing/opentracing-go/mocktracer

     1  package mocktracer
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/opentracing/opentracing-go"
     7  )
     8  
     9  // New returns a MockTracer opentracing.Tracer implementation that's intended
    10  // to facilitate tests of OpenTracing instrumentation.
    11  func New() *MockTracer {
    12  	t := &MockTracer{
    13  		finishedSpans: []*MockSpan{},
    14  		startedSpans:  []*MockSpan{},
    15  		injectors:     make(map[interface{}]Injector),
    16  		extractors:    make(map[interface{}]Extractor),
    17  	}
    18  
    19  	// register default injectors/extractors
    20  	textPropagator := new(TextMapPropagator)
    21  	t.RegisterInjector(opentracing.TextMap, textPropagator)
    22  	t.RegisterExtractor(opentracing.TextMap, textPropagator)
    23  
    24  	httpPropagator := &TextMapPropagator{HTTPHeaders: true}
    25  	t.RegisterInjector(opentracing.HTTPHeaders, httpPropagator)
    26  	t.RegisterExtractor(opentracing.HTTPHeaders, httpPropagator)
    27  
    28  	return t
    29  }
    30  
    31  // MockTracer is only intended for testing OpenTracing instrumentation.
    32  //
    33  // It is entirely unsuitable for production use, but appropriate for tests
    34  // that want to verify tracing behavior in other frameworks/applications.
    35  type MockTracer struct {
    36  	sync.RWMutex
    37  	finishedSpans []*MockSpan
    38  	startedSpans  []*MockSpan
    39  	injectors     map[interface{}]Injector
    40  	extractors    map[interface{}]Extractor
    41  }
    42  
    43  // UnfinishedSpans returns all spans that have been started and not finished since the
    44  // MockTracer was constructed or since the last call to its Reset() method.
    45  func (t *MockTracer) UnfinishedSpans() []*MockSpan {
    46  	t.RLock()
    47  	defer t.RUnlock()
    48  
    49  	spans := make([]*MockSpan, len(t.startedSpans))
    50  	copy(spans, t.startedSpans)
    51  	return spans
    52  }
    53  
    54  // FinishedSpans returns all spans that have been Finish()'ed since the
    55  // MockTracer was constructed or since the last call to its Reset() method.
    56  func (t *MockTracer) FinishedSpans() []*MockSpan {
    57  	t.RLock()
    58  	defer t.RUnlock()
    59  	spans := make([]*MockSpan, len(t.finishedSpans))
    60  	copy(spans, t.finishedSpans)
    61  	return spans
    62  }
    63  
    64  // Reset clears the internally accumulated finished spans. Note that any
    65  // extant MockSpans will still append to finishedSpans when they Finish(),
    66  // even after a call to Reset().
    67  func (t *MockTracer) Reset() {
    68  	t.Lock()
    69  	defer t.Unlock()
    70  	t.startedSpans = []*MockSpan{}
    71  	t.finishedSpans = []*MockSpan{}
    72  }
    73  
    74  // StartSpan belongs to the Tracer interface.
    75  func (t *MockTracer) StartSpan(operationName string, opts ...opentracing.StartSpanOption) opentracing.Span {
    76  	sso := opentracing.StartSpanOptions{}
    77  	for _, o := range opts {
    78  		o.Apply(&sso)
    79  	}
    80  
    81  	span := newMockSpan(t, operationName, sso)
    82  	t.recordStartedSpan(span)
    83  	return span
    84  }
    85  
    86  // RegisterInjector registers injector for given format
    87  func (t *MockTracer) RegisterInjector(format interface{}, injector Injector) {
    88  	t.injectors[format] = injector
    89  }
    90  
    91  // RegisterExtractor registers extractor for given format
    92  func (t *MockTracer) RegisterExtractor(format interface{}, extractor Extractor) {
    93  	t.extractors[format] = extractor
    94  }
    95  
    96  // Inject belongs to the Tracer interface.
    97  func (t *MockTracer) Inject(sm opentracing.SpanContext, format interface{}, carrier interface{}) error {
    98  	spanContext, ok := sm.(MockSpanContext)
    99  	if !ok {
   100  		return opentracing.ErrInvalidSpanContext
   101  	}
   102  	injector, ok := t.injectors[format]
   103  	if !ok {
   104  		return opentracing.ErrUnsupportedFormat
   105  	}
   106  	return injector.Inject(spanContext, carrier)
   107  }
   108  
   109  // Extract belongs to the Tracer interface.
   110  func (t *MockTracer) Extract(format interface{}, carrier interface{}) (opentracing.SpanContext, error) {
   111  	extractor, ok := t.extractors[format]
   112  	if !ok {
   113  		return nil, opentracing.ErrUnsupportedFormat
   114  	}
   115  	return extractor.Extract(carrier)
   116  }
   117  
   118  func (t *MockTracer) recordStartedSpan(span *MockSpan) {
   119  	t.Lock()
   120  	defer t.Unlock()
   121  	t.startedSpans = append(t.startedSpans, span)
   122  }
   123  
   124  func (t *MockTracer) recordFinishedSpan(span *MockSpan) {
   125  	t.Lock()
   126  	defer t.Unlock()
   127  	t.finishedSpans = append(t.finishedSpans, span)
   128  
   129  	for i := range t.startedSpans {
   130  		if t.startedSpans[i].SpanContext.SpanID == span.SpanContext.SpanID &&
   131  			t.startedSpans[i].SpanContext.TraceID == span.SpanContext.TraceID {
   132  			t.startedSpans = append(t.startedSpans[:i], t.startedSpans[i+1:]...)
   133  			return
   134  		}
   135  	}
   136  }
   137
View as plain text