1package mocktracer
2
3import (
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.
11func New() *MockTracer {
12	t := &MockTracer{
13		finishedSpans: []*MockSpan{},
14		injectors:     make(map[interface{}]Injector),
15		extractors:    make(map[interface{}]Extractor),
16	}
17
18	// register default injectors/extractors
19	textPropagator := new(TextMapPropagator)
20	t.RegisterInjector(opentracing.TextMap, textPropagator)
21	t.RegisterExtractor(opentracing.TextMap, textPropagator)
22
23	httpPropagator := &TextMapPropagator{HTTPHeaders: true}
24	t.RegisterInjector(opentracing.HTTPHeaders, httpPropagator)
25	t.RegisterExtractor(opentracing.HTTPHeaders, httpPropagator)
26
27	return t
28}
29
30// MockTracer is only intended for testing OpenTracing instrumentation.
31//
32// It is entirely unsuitable for production use, but appropriate for tests
33// that want to verify tracing behavior in other frameworks/applications.
34type MockTracer struct {
35	sync.RWMutex
36	finishedSpans []*MockSpan
37	injectors     map[interface{}]Injector
38	extractors    map[interface{}]Extractor
39}
40
41// FinishedSpans returns all spans that have been Finish()'ed since the
42// MockTracer was constructed or since the last call to its Reset() method.
43func (t *MockTracer) FinishedSpans() []*MockSpan {
44	t.RLock()
45	defer t.RUnlock()
46	spans := make([]*MockSpan, len(t.finishedSpans))
47	copy(spans, t.finishedSpans)
48	return spans
49}
50
51// Reset clears the internally accumulated finished spans. Note that any
52// extant MockSpans will still append to finishedSpans when they Finish(),
53// even after a call to Reset().
54func (t *MockTracer) Reset() {
55	t.Lock()
56	defer t.Unlock()
57	t.finishedSpans = []*MockSpan{}
58}
59
60// StartSpan belongs to the Tracer interface.
61func (t *MockTracer) StartSpan(operationName string, opts ...opentracing.StartSpanOption) opentracing.Span {
62	sso := opentracing.StartSpanOptions{}
63	for _, o := range opts {
64		o.Apply(&sso)
65	}
66	return newMockSpan(t, operationName, sso)
67}
68
69// RegisterInjector registers injector for given format
70func (t *MockTracer) RegisterInjector(format interface{}, injector Injector) {
71	t.injectors[format] = injector
72}
73
74// RegisterExtractor registers extractor for given format
75func (t *MockTracer) RegisterExtractor(format interface{}, extractor Extractor) {
76	t.extractors[format] = extractor
77}
78
79// Inject belongs to the Tracer interface.
80func (t *MockTracer) Inject(sm opentracing.SpanContext, format interface{}, carrier interface{}) error {
81	spanContext, ok := sm.(MockSpanContext)
82	if !ok {
83		return opentracing.ErrInvalidSpanContext
84	}
85	injector, ok := t.injectors[format]
86	if !ok {
87		return opentracing.ErrUnsupportedFormat
88	}
89	return injector.Inject(spanContext, carrier)
90}
91
92// Extract belongs to the Tracer interface.
93func (t *MockTracer) Extract(format interface{}, carrier interface{}) (opentracing.SpanContext, error) {
94	extractor, ok := t.extractors[format]
95	if !ok {
96		return nil, opentracing.ErrUnsupportedFormat
97	}
98	return extractor.Extract(carrier)
99}
100
101func (t *MockTracer) recordSpan(span *MockSpan) {
102	t.Lock()
103	defer t.Unlock()
104	t.finishedSpans = append(t.finishedSpans, span)
105}
106