mockspan.go

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

     1  package mocktracer
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"sync/atomic"
     7  	"time"
     8  
     9  	"github.com/opentracing/opentracing-go"
    10  	"github.com/opentracing/opentracing-go/ext"
    11  	"github.com/opentracing/opentracing-go/log"
    12  )
    13  
    14  // MockSpanContext is an opentracing.SpanContext implementation.
    15  //
    16  // It is entirely unsuitable for production use, but appropriate for tests
    17  // that want to verify tracing behavior in other frameworks/applications.
    18  //
    19  // By default all spans have Sampled=true flag, unless {"sampling.priority": 0}
    20  // tag is set.
    21  type MockSpanContext struct {
    22  	TraceID int
    23  	SpanID  int
    24  	Sampled bool
    25  	Baggage map[string]string
    26  }
    27  
    28  var mockIDSource = uint32(42)
    29  
    30  func nextMockID() int {
    31  	return int(atomic.AddUint32(&mockIDSource, 1))
    32  }
    33  
    34  // ForeachBaggageItem belongs to the SpanContext interface
    35  func (c MockSpanContext) ForeachBaggageItem(handler func(k, v string) bool) {
    36  	for k, v := range c.Baggage {
    37  		if !handler(k, v) {
    38  			break
    39  		}
    40  	}
    41  }
    42  
    43  // WithBaggageItem creates a new context with an extra baggage item.
    44  func (c MockSpanContext) WithBaggageItem(key, value string) MockSpanContext {
    45  	var newBaggage map[string]string
    46  	if c.Baggage == nil {
    47  		newBaggage = map[string]string{key: value}
    48  	} else {
    49  		newBaggage = make(map[string]string, len(c.Baggage)+1)
    50  		for k, v := range c.Baggage {
    51  			newBaggage[k] = v
    52  		}
    53  		newBaggage[key] = value
    54  	}
    55  	// Use positional parameters so the compiler will help catch new fields.
    56  	return MockSpanContext{c.TraceID, c.SpanID, c.Sampled, newBaggage}
    57  }
    58  
    59  // MockSpan is an opentracing.Span implementation that exports its internal
    60  // state for testing purposes.
    61  type MockSpan struct {
    62  	sync.RWMutex
    63  
    64  	ParentID int
    65  
    66  	OperationName string
    67  	StartTime     time.Time
    68  	FinishTime    time.Time
    69  
    70  	// All of the below are protected by the embedded RWMutex.
    71  	SpanContext MockSpanContext
    72  	tags        map[string]interface{}
    73  	logs        []MockLogRecord
    74  	tracer      *MockTracer
    75  }
    76  
    77  func newMockSpan(t *MockTracer, name string, opts opentracing.StartSpanOptions) *MockSpan {
    78  	tags := opts.Tags
    79  	if tags == nil {
    80  		tags = map[string]interface{}{}
    81  	}
    82  	traceID := nextMockID()
    83  	parentID := int(0)
    84  	var baggage map[string]string
    85  	sampled := true
    86  	if len(opts.References) > 0 {
    87  		traceID = opts.References[0].ReferencedContext.(MockSpanContext).TraceID
    88  		parentID = opts.References[0].ReferencedContext.(MockSpanContext).SpanID
    89  		sampled = opts.References[0].ReferencedContext.(MockSpanContext).Sampled
    90  		baggage = opts.References[0].ReferencedContext.(MockSpanContext).Baggage
    91  	}
    92  	spanContext := MockSpanContext{traceID, nextMockID(), sampled, baggage}
    93  	startTime := opts.StartTime
    94  	if startTime.IsZero() {
    95  		startTime = time.Now()
    96  	}
    97  	return &MockSpan{
    98  		ParentID:      parentID,
    99  		OperationName: name,
   100  		StartTime:     startTime,
   101  		tags:          tags,
   102  		logs:          []MockLogRecord{},
   103  		SpanContext:   spanContext,
   104  
   105  		tracer: t,
   106  	}
   107  }
   108  
   109  // Tags returns a copy of tags accumulated by the span so far
   110  func (s *MockSpan) Tags() map[string]interface{} {
   111  	s.RLock()
   112  	defer s.RUnlock()
   113  	tags := make(map[string]interface{})
   114  	for k, v := range s.tags {
   115  		tags[k] = v
   116  	}
   117  	return tags
   118  }
   119  
   120  // Tag returns a single tag
   121  func (s *MockSpan) Tag(k string) interface{} {
   122  	s.RLock()
   123  	defer s.RUnlock()
   124  	return s.tags[k]
   125  }
   126  
   127  // Logs returns a copy of logs accumulated in the span so far
   128  func (s *MockSpan) Logs() []MockLogRecord {
   129  	s.RLock()
   130  	defer s.RUnlock()
   131  	logs := make([]MockLogRecord, len(s.logs))
   132  	copy(logs, s.logs)
   133  	return logs
   134  }
   135  
   136  // Context belongs to the Span interface
   137  func (s *MockSpan) Context() opentracing.SpanContext {
   138  	s.Lock()
   139  	defer s.Unlock()
   140  	return s.SpanContext
   141  }
   142  
   143  // SetTag belongs to the Span interface
   144  func (s *MockSpan) SetTag(key string, value interface{}) opentracing.Span {
   145  	s.Lock()
   146  	defer s.Unlock()
   147  	if key == string(ext.SamplingPriority) {
   148  		if v, ok := value.(uint16); ok {
   149  			s.SpanContext.Sampled = v > 0
   150  			return s
   151  		}
   152  		if v, ok := value.(int); ok {
   153  			s.SpanContext.Sampled = v > 0
   154  			return s
   155  		}
   156  	}
   157  	s.tags[key] = value
   158  	return s
   159  }
   160  
   161  // SetBaggageItem belongs to the Span interface
   162  func (s *MockSpan) SetBaggageItem(key, val string) opentracing.Span {
   163  	s.Lock()
   164  	defer s.Unlock()
   165  	s.SpanContext = s.SpanContext.WithBaggageItem(key, val)
   166  	return s
   167  }
   168  
   169  // BaggageItem belongs to the Span interface
   170  func (s *MockSpan) BaggageItem(key string) string {
   171  	s.RLock()
   172  	defer s.RUnlock()
   173  	return s.SpanContext.Baggage[key]
   174  }
   175  
   176  // Finish belongs to the Span interface
   177  func (s *MockSpan) Finish() {
   178  	s.Lock()
   179  	s.FinishTime = time.Now()
   180  	s.Unlock()
   181  	s.tracer.recordFinishedSpan(s)
   182  }
   183  
   184  // FinishWithOptions belongs to the Span interface
   185  func (s *MockSpan) FinishWithOptions(opts opentracing.FinishOptions) {
   186  	s.Lock()
   187  	s.FinishTime = opts.FinishTime
   188  	s.Unlock()
   189  
   190  	// Handle any late-bound LogRecords.
   191  	for _, lr := range opts.LogRecords {
   192  		s.logFieldsWithTimestamp(lr.Timestamp, lr.Fields...)
   193  	}
   194  	// Handle (deprecated) BulkLogData.
   195  	for _, ld := range opts.BulkLogData {
   196  		if ld.Payload != nil {
   197  			s.logFieldsWithTimestamp(
   198  				ld.Timestamp,
   199  				log.String("event", ld.Event),
   200  				log.Object("payload", ld.Payload))
   201  		} else {
   202  			s.logFieldsWithTimestamp(
   203  				ld.Timestamp,
   204  				log.String("event", ld.Event))
   205  		}
   206  	}
   207  
   208  	s.tracer.recordFinishedSpan(s)
   209  }
   210  
   211  // String allows printing span for debugging
   212  func (s *MockSpan) String() string {
   213  	return fmt.Sprintf(
   214  		"traceId=%d, spanId=%d, parentId=%d, sampled=%t, name=%s",
   215  		s.SpanContext.TraceID, s.SpanContext.SpanID, s.ParentID,
   216  		s.SpanContext.Sampled, s.OperationName)
   217  }
   218  
   219  // LogFields belongs to the Span interface
   220  func (s *MockSpan) LogFields(fields ...log.Field) {
   221  	s.logFieldsWithTimestamp(time.Now(), fields...)
   222  }
   223  
   224  // The caller MUST NOT hold s.Lock
   225  func (s *MockSpan) logFieldsWithTimestamp(ts time.Time, fields ...log.Field) {
   226  	lr := MockLogRecord{
   227  		Timestamp: ts,
   228  		Fields:    make([]MockKeyValue, len(fields)),
   229  	}
   230  	for i, f := range fields {
   231  		outField := &(lr.Fields[i])
   232  		f.Marshal(outField)
   233  	}
   234  
   235  	s.Lock()
   236  	defer s.Unlock()
   237  	s.logs = append(s.logs, lr)
   238  }
   239  
   240  // LogKV belongs to the Span interface.
   241  //
   242  // This implementations coerces all "values" to strings, though that is not
   243  // something all implementations need to do. Indeed, a motivated person can and
   244  // probably should have this do a typed switch on the values.
   245  func (s *MockSpan) LogKV(keyValues ...interface{}) {
   246  	if len(keyValues)%2 != 0 {
   247  		s.LogFields(log.Error(fmt.Errorf("Non-even keyValues len: %v", len(keyValues))))
   248  		return
   249  	}
   250  	fields, err := log.InterleavedKVToFields(keyValues...)
   251  	if err != nil {
   252  		s.LogFields(log.Error(err), log.String("function", "LogKV"))
   253  		return
   254  	}
   255  	s.LogFields(fields...)
   256  }
   257  
   258  // LogEvent belongs to the Span interface
   259  func (s *MockSpan) LogEvent(event string) {
   260  	s.LogFields(log.String("event", event))
   261  }
   262  
   263  // LogEventWithPayload belongs to the Span interface
   264  func (s *MockSpan) LogEventWithPayload(event string, payload interface{}) {
   265  	s.LogFields(log.String("event", event), log.Object("payload", payload))
   266  }
   267  
   268  // Log belongs to the Span interface
   269  func (s *MockSpan) Log(data opentracing.LogData) {
   270  	panic("MockSpan.Log() no longer supported")
   271  }
   272  
   273  // SetOperationName belongs to the Span interface
   274  func (s *MockSpan) SetOperationName(operationName string) opentracing.Span {
   275  	s.Lock()
   276  	defer s.Unlock()
   277  	s.OperationName = operationName
   278  	return s
   279  }
   280  
   281  // Tracer belongs to the Span interface
   282  func (s *MockSpan) Tracer() opentracing.Tracer {
   283  	return s.tracer
   284  }
   285
View as plain text