ipld-eth-server/vendor/gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer/mocktracer.go
2018-09-11 16:30:29 -05:00

160 lines
4.2 KiB
Go

// Package mocktracer provides a mock implementation of the tracer used in testing. It
// allows querying spans generated at runtime, without having them actually be sent to
// an agent. It provides a simple way to test that instrumentation is running correctly
// in your application.
//
// Simply call "Start" at the beginning of your tests to start and obtain an instance
// of the mock tracer.
package mocktracer
import (
"strconv"
"strings"
"sync"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)
var _ ddtrace.Tracer = (*mocktracer)(nil)
var _ Tracer = (*mocktracer)(nil)
// Tracer exposes an interface for querying the currently running mock tracer.
type Tracer interface {
// FinishedSpans returns the set of finished spans.
FinishedSpans() []Span
// Reset resets the spans and services recorded in the tracer. This is
// especially useful when running tests in a loop, where a clean start
// is desired for FinishedSpans calls.
Reset()
// Stop deactivates the mock tracer and allows a normal tracer to take over.
// It should always be called when testing has finished.
Stop()
}
// Start sets the internal tracer to a mock and returns an interface
// which allows querying it. Call Start at the beginning of your tests
// to activate the mock tracer. When your test runs, use the returned
// interface to query the tracer's state.
func Start() Tracer {
var t mocktracer
internal.SetGlobalTracer(&t)
internal.Testing = true
return &t
}
type mocktracer struct {
sync.RWMutex // guards below spans
finishedSpans []Span
}
// Stop deactivates the mock tracer and sets the active tracer to a no-op.
func (*mocktracer) Stop() {
internal.SetGlobalTracer(&internal.NoopTracer{})
internal.Testing = false
}
func (t *mocktracer) StartSpan(operationName string, opts ...ddtrace.StartSpanOption) ddtrace.Span {
var cfg ddtrace.StartSpanConfig
for _, fn := range opts {
fn(&cfg)
}
return newSpan(t, operationName, &cfg)
}
func (t *mocktracer) FinishedSpans() []Span {
t.RLock()
defer t.RUnlock()
return t.finishedSpans
}
func (t *mocktracer) Reset() {
t.Lock()
defer t.Unlock()
t.finishedSpans = nil
}
func (t *mocktracer) addFinishedSpan(s Span) {
t.Lock()
defer t.Unlock()
if t.finishedSpans == nil {
t.finishedSpans = make([]Span, 0, 1)
}
t.finishedSpans = append(t.finishedSpans, s)
}
const (
traceHeader = tracer.DefaultTraceIDHeader
spanHeader = tracer.DefaultParentIDHeader
priorityHeader = tracer.DefaultPriorityHeader
baggagePrefix = tracer.DefaultBaggageHeaderPrefix
)
func (t *mocktracer) Extract(carrier interface{}) (ddtrace.SpanContext, error) {
reader, ok := carrier.(tracer.TextMapReader)
if !ok {
return nil, tracer.ErrInvalidCarrier
}
var sc spanContext
err := reader.ForeachKey(func(key, v string) error {
k := strings.ToLower(key)
if k == traceHeader {
id, err := strconv.ParseUint(v, 10, 64)
if err != nil {
return tracer.ErrSpanContextCorrupted
}
sc.traceID = id
}
if k == spanHeader {
id, err := strconv.ParseUint(v, 10, 64)
if err != nil {
return tracer.ErrSpanContextCorrupted
}
sc.spanID = id
}
if k == priorityHeader {
p, err := strconv.Atoi(v)
if err != nil {
return tracer.ErrSpanContextCorrupted
}
sc.priority = p
sc.hasPriority = true
}
if strings.HasPrefix(k, baggagePrefix) {
sc.setBaggageItem(strings.TrimPrefix(k, baggagePrefix), v)
}
return nil
})
if err != nil {
return nil, err
}
if sc.traceID == 0 || sc.spanID == 0 {
return nil, tracer.ErrSpanContextNotFound
}
return &sc, err
}
func (t *mocktracer) Inject(context ddtrace.SpanContext, carrier interface{}) error {
writer, ok := carrier.(tracer.TextMapWriter)
if !ok {
return tracer.ErrInvalidCarrier
}
ctx, ok := context.(*spanContext)
if !ok || ctx.traceID == 0 || ctx.spanID == 0 {
return tracer.ErrInvalidSpanContext
}
writer.Set(traceHeader, strconv.FormatUint(ctx.traceID, 10))
writer.Set(spanHeader, strconv.FormatUint(ctx.spanID, 10))
if ctx.hasSamplingPriority() {
writer.Set(priorityHeader, strconv.Itoa(ctx.priority))
}
ctx.ForeachBaggageItem(func(k, v string) bool {
writer.Set(baggagePrefix+k, v)
return true
})
return nil
}