171 lines
4.4 KiB
Go
171 lines
4.4 KiB
Go
|
package tracer
|
||
|
|
||
|
import (
|
||
|
"errors"
|
||
|
"net/http"
|
||
|
"strconv"
|
||
|
"testing"
|
||
|
|
||
|
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
|
||
|
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
|
||
|
|
||
|
"github.com/stretchr/testify/assert"
|
||
|
)
|
||
|
|
||
|
func TestHTTPHeadersCarrierSet(t *testing.T) {
|
||
|
h := http.Header{}
|
||
|
c := HTTPHeadersCarrier(h)
|
||
|
c.Set("A", "x")
|
||
|
assert.Equal(t, "x", h.Get("A"))
|
||
|
}
|
||
|
|
||
|
func TestHTTPHeadersCarrierForeachKey(t *testing.T) {
|
||
|
h := http.Header{}
|
||
|
h.Add("A", "x")
|
||
|
h.Add("B", "y")
|
||
|
got := map[string]string{}
|
||
|
err := HTTPHeadersCarrier(h).ForeachKey(func(k, v string) error {
|
||
|
got[k] = v
|
||
|
return nil
|
||
|
})
|
||
|
assert := assert.New(t)
|
||
|
assert.Nil(err)
|
||
|
assert.Equal("x", h.Get("A"))
|
||
|
assert.Equal("y", h.Get("B"))
|
||
|
}
|
||
|
|
||
|
func TestHTTPHeadersCarrierForeachKeyError(t *testing.T) {
|
||
|
want := errors.New("random error")
|
||
|
h := http.Header{}
|
||
|
h.Add("A", "x")
|
||
|
h.Add("B", "y")
|
||
|
got := HTTPHeadersCarrier(h).ForeachKey(func(k, v string) error {
|
||
|
if k == "B" {
|
||
|
return want
|
||
|
}
|
||
|
return nil
|
||
|
})
|
||
|
assert.Equal(t, want, got)
|
||
|
}
|
||
|
|
||
|
func TestTextMapCarrierSet(t *testing.T) {
|
||
|
m := map[string]string{}
|
||
|
c := TextMapCarrier(m)
|
||
|
c.Set("a", "b")
|
||
|
assert.Equal(t, "b", m["a"])
|
||
|
}
|
||
|
|
||
|
func TestTextMapCarrierForeachKey(t *testing.T) {
|
||
|
want := map[string]string{"A": "x", "B": "y"}
|
||
|
got := map[string]string{}
|
||
|
err := TextMapCarrier(want).ForeachKey(func(k, v string) error {
|
||
|
got[k] = v
|
||
|
return nil
|
||
|
})
|
||
|
assert := assert.New(t)
|
||
|
assert.Nil(err)
|
||
|
assert.Equal(got, want)
|
||
|
}
|
||
|
|
||
|
func TestTextMapCarrierForeachKeyError(t *testing.T) {
|
||
|
m := map[string]string{"A": "x", "B": "y"}
|
||
|
want := errors.New("random error")
|
||
|
got := TextMapCarrier(m).ForeachKey(func(k, v string) error {
|
||
|
return want
|
||
|
})
|
||
|
assert.Equal(t, got, want)
|
||
|
}
|
||
|
|
||
|
func TestTextMapPropagatorErrors(t *testing.T) {
|
||
|
propagator := NewPropagator(nil)
|
||
|
assert := assert.New(t)
|
||
|
|
||
|
err := propagator.Inject(&spanContext{}, 2)
|
||
|
assert.Equal(ErrInvalidCarrier, err)
|
||
|
err = propagator.Inject(internal.NoopSpanContext{}, TextMapCarrier(map[string]string{}))
|
||
|
assert.Equal(ErrInvalidSpanContext, err)
|
||
|
err = propagator.Inject(&spanContext{}, TextMapCarrier(map[string]string{}))
|
||
|
assert.Equal(ErrInvalidSpanContext, err) // no traceID and spanID
|
||
|
err = propagator.Inject(&spanContext{traceID: 1}, TextMapCarrier(map[string]string{}))
|
||
|
assert.Equal(ErrInvalidSpanContext, err) // no spanID
|
||
|
|
||
|
_, err = propagator.Extract(2)
|
||
|
assert.Equal(ErrInvalidCarrier, err)
|
||
|
|
||
|
_, err = propagator.Extract(TextMapCarrier(map[string]string{
|
||
|
DefaultTraceIDHeader: "1",
|
||
|
DefaultParentIDHeader: "A",
|
||
|
}))
|
||
|
assert.Equal(ErrSpanContextCorrupted, err)
|
||
|
|
||
|
_, err = propagator.Extract(TextMapCarrier(map[string]string{
|
||
|
DefaultTraceIDHeader: "A",
|
||
|
DefaultParentIDHeader: "2",
|
||
|
}))
|
||
|
assert.Equal(ErrSpanContextCorrupted, err)
|
||
|
|
||
|
_, err = propagator.Extract(TextMapCarrier(map[string]string{
|
||
|
DefaultTraceIDHeader: "0",
|
||
|
DefaultParentIDHeader: "0",
|
||
|
}))
|
||
|
assert.Equal(ErrSpanContextNotFound, err)
|
||
|
}
|
||
|
|
||
|
func TestTextMapPropagatorInjectHeader(t *testing.T) {
|
||
|
assert := assert.New(t)
|
||
|
|
||
|
propagator := NewPropagator(&PropagatorConfig{
|
||
|
BaggagePrefix: "bg-",
|
||
|
TraceHeader: "tid",
|
||
|
ParentHeader: "pid",
|
||
|
})
|
||
|
tracer := newTracer(WithPropagator(propagator))
|
||
|
|
||
|
root := tracer.StartSpan("web.request").(*span)
|
||
|
root.SetBaggageItem("item", "x")
|
||
|
root.SetTag(ext.SamplingPriority, 0)
|
||
|
ctx := root.Context()
|
||
|
headers := http.Header{}
|
||
|
|
||
|
carrier := HTTPHeadersCarrier(headers)
|
||
|
err := tracer.Inject(ctx, carrier)
|
||
|
assert.Nil(err)
|
||
|
|
||
|
tid := strconv.FormatUint(root.TraceID, 10)
|
||
|
pid := strconv.FormatUint(root.SpanID, 10)
|
||
|
|
||
|
assert.Equal(headers.Get("tid"), tid)
|
||
|
assert.Equal(headers.Get("pid"), pid)
|
||
|
assert.Equal(headers.Get("bg-item"), "x")
|
||
|
assert.Equal(headers.Get(DefaultPriorityHeader), "0")
|
||
|
}
|
||
|
|
||
|
func TestTextMapPropagatorInjectExtract(t *testing.T) {
|
||
|
propagator := NewPropagator(&PropagatorConfig{
|
||
|
BaggagePrefix: "bg-",
|
||
|
TraceHeader: "tid",
|
||
|
ParentHeader: "pid",
|
||
|
})
|
||
|
tracer := newTracer(WithPropagator(propagator))
|
||
|
root := tracer.StartSpan("web.request").(*span)
|
||
|
root.SetTag(ext.SamplingPriority, -1)
|
||
|
root.SetBaggageItem("item", "x")
|
||
|
ctx := root.Context().(*spanContext)
|
||
|
headers := TextMapCarrier(map[string]string{})
|
||
|
err := tracer.Inject(ctx, headers)
|
||
|
|
||
|
assert := assert.New(t)
|
||
|
assert.Nil(err)
|
||
|
|
||
|
sctx, err := tracer.Extract(headers)
|
||
|
assert.Nil(err)
|
||
|
|
||
|
xctx, ok := sctx.(*spanContext)
|
||
|
assert.True(ok)
|
||
|
assert.Equal(xctx.traceID, ctx.traceID)
|
||
|
assert.Equal(xctx.spanID, ctx.spanID)
|
||
|
assert.Equal(xctx.baggage, ctx.baggage)
|
||
|
assert.Equal(xctx.priority, ctx.priority)
|
||
|
assert.Equal(xctx.hasPriority, ctx.hasPriority)
|
||
|
}
|