160 lines
4.2 KiB
Go
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
|
|
}
|