Add unit tests (#203)

* B-30111 - Add UTs

* B-30111 - Add UTs

* Add unit tests

tracing/annotator.go

@@ -17,7 +17,7 @@ const (
 var defaultFormat propagation.HTTPFormat = &b3.HTTPFormat{}
 
 //SpanContextAnnotator retrieve information about current span from context or HTTP headers
-//and propogate in binary format to gRPC service
+//and propagate in binary format to gRPC service
 func SpanContextAnnotator(ctx context.Context, req *http.Request) metadata.MD {
 	md := make(metadata.MD)
 

tracing/annotator_test.go

@@ -0,0 +1,28 @@
+package tracing
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.opencensus.io/trace"
+	"go.opencensus.io/trace/propagation"
+)
+
+func TestSpanContextAnnotator_FromContext(t *testing.T) {
+	ctx, span := trace.StartSpan(context.Background(), "")
+	result := SpanContextAnnotator(ctx, nil)
+	assert.Equal(t, []string{string(propagation.Binary(span.SpanContext()))}, result[traceContextKey])
+}
+
+func TestSpanContextAnnotator_FromRequest(t *testing.T) {
+	_, span := trace.StartSpan(context.Background(), "")
+	req, _ := http.NewRequest("", "", nil)
+	defaultFormat.SpanContextToRequest(span.SpanContext(), req)
+	sc, ok := defaultFormat.SpanContextFromRequest(req)
+	assert.True(t, ok)
+
+	result := SpanContextAnnotator(context.Background(), req)
+	assert.Equal(t, []string{string(propagation.Binary(sc))}, result[traceContextKey])
+}

tracing/grpc.go

@@ -123,7 +123,7 @@ func (s *ServerHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 			attrs := metadataToAttributes(rs.Trailer, RequestTrailerAnnotationPrefix, s.options.metadataMatcher)
 			span.AddAttributes(attrs...)
 		case *stats.OutHeader:
-			attrs := metadataToAttributes(rs.Header, ResponsePayloadAnnotationKey, s.options.metadataMatcher)
+			attrs := metadataToAttributes(rs.Header, ResponseHeaderAnnotationPrefix, s.options.metadataMatcher)
 			span.AddAttributes(attrs...)
 		case *stats.OutTrailer:
 			attrs := metadataToAttributes(rs.Trailer, ResponseTrailerAnnotationPrefix, s.options.metadataMatcher)
@@ -162,10 +162,10 @@ func (s *ServerHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 
 }
 
-func metadataToAttributes(md metadata.MD, prefix string, marcher metadataMatcher) []trace.Attribute {
+func metadataToAttributes(md metadata.MD, prefix string, matcher metadataMatcher) []trace.Attribute {
 	attrs := make([]trace.Attribute, 0, len(md))
 	for k, vals := range md {
-		key, ok := marcher(k)
+		key, ok := matcher(k)
 		if !ok {
 			continue
 		}

tracing/grpc_test.go

@@ -0,0 +1,180 @@
+package tracing
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.opencensus.io/trace"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
+)
+
+var testGRPCOptions = &gRPCOptions{}
+
+func TestDefaultGRPCOptions(t *testing.T) {
+	expected := &gRPCOptions{
+		metadataMatcher: defaultMetadataMatcher,
+		maxPayloadSize:  1048576,
+	}
+
+	result := defaultGRPCOptions()
+	expectedHeader, expectedBool := expected.metadataMatcher(expectedStr)
+	resultHeader, resultBool := result.metadataMatcher(expectedStr)
+	assert.True(t, expectedBool)
+	assert.Equal(t, expectedBool, resultBool)
+	assert.Equal(t, expectedHeader, resultHeader)
+	assert.Equal(t, expected.maxPayloadSize, result.maxPayloadSize)
+}
+
+func TestWithMetadataAnnotation(t *testing.T) {
+	option := WithMetadataAnnotation(func(ctx context.Context, stats stats.RPCStats) bool {
+		return true
+	})
+	option(testGRPCOptions)
+	assert.True(t, testGRPCOptions.spanWithMetadata(nil, nil))
+}
+
+func TestWithMetadataMatcher(t *testing.T) {
+	option := WithMetadataMatcher(func(s string) (string, bool) {
+		return s, true
+	})
+	option(testGRPCOptions)
+	resultStr, ok := testGRPCOptions.metadataMatcher(expectedStr)
+	assert.True(t, ok)
+	assert.Equal(t, expectedStr, resultStr)
+}
+
+func TestWithGRPCPayloadAnnotation(t *testing.T) {
+	option := WithGRPCPayloadAnnotation(func(ctx context.Context, rpcStats stats.RPCStats) bool {
+		return true
+	})
+	option(testGRPCOptions)
+	assert.True(t, testGRPCOptions.spanWithPayload(nil, nil))
+}
+
+func TestWithGRPCPayloadLimit(t *testing.T) {
+	option := WithGRPCPayloadLimit(333)
+	option(testGRPCOptions)
+	assert.Equal(t, 333, testGRPCOptions.maxPayloadSize)
+}
+
+func TestNewServerHandler(t *testing.T) {
+	result := NewServerHandler(func(options *gRPCOptions) {
+		options.spanWithPayload = func(ctx context.Context, rpcStats stats.RPCStats) bool {
+			return true
+		}
+	})
+
+	matcherStr, ok := result.options.metadataMatcher(expectedStr)
+	assert.True(t, ok)
+	assert.Equal(t, expectedStr, matcherStr)
+	assert.True(t, result.options.spanWithPayload(nil, nil))
+	assert.Equal(t, DefaultMaxPayloadSize, result.options.maxPayloadSize)
+}
+
+func TestServerHandler_HandleRPC(t *testing.T) {
+	handler := NewServerHandler(func(options *gRPCOptions) {
+		options.spanWithPayload = func(ctx context.Context, rpcStats stats.RPCStats) bool {
+			return true
+		}
+
+		options.spanWithMetadata = func(ctx context.Context, rpcStats stats.RPCStats) bool {
+			return true
+		}
+	})
+
+	expectedStats := []stats.RPCStats{
+		&stats.End{
+			Error: fmt.Errorf(""),
+		},
+		&stats.InHeader{
+			Header: map[string][]string{
+				"header1": {""},
+			},
+		},
+		&stats.InTrailer{
+			Trailer: map[string][]string{
+				"trailer1": {""},
+			},
+		},
+		&stats.OutHeader{
+			Header: map[string][]string{
+				"outHeader1": {""},
+			},
+		},
+		&stats.OutTrailer{
+			Trailer: map[string][]string{
+				"outTrailer1": {""},
+			},
+		},
+		&stats.InPayload{
+			Payload: []byte(""),
+		},
+		&stats.OutPayload{
+			Payload: []byte(""),
+		},
+	}
+
+	ctx, _ := trace.StartSpan(context.Background(), "test span", trace.WithSampler(trace.AlwaysSample()))
+
+	for _, v := range expectedStats {
+		handler.HandleRPC(ctx, v)
+	}
+
+	expectedMap := map[string]string{
+		"request.header.header1":       "true",
+		"request.trailer.trailer1":     "true",
+		"response.header.outHeader1":   "true",
+		"response.trailer.outTrailer1": "true",
+	}
+
+	resultMap := make(map[string]string, 4)
+	reflectAttrs := reflect.ValueOf(trace.FromContext(ctx)).Elem().Field(3).Elem().Field(0)
+	reflectKeys := reflectAttrs.MapKeys()
+	for _, k := range reflectKeys {
+		key := k.Convert(reflectAttrs.Type().Key())
+		val := reflectAttrs.MapIndex(key)
+		resultMap[fmt.Sprint(key)] = fmt.Sprint(val)
+	}
+
+	assert.Equal(t, expectedMap, resultMap)
+
+	expectedAnnotations := []string{
+		"Response error", "Request payload", "Response payload",
+	}
+
+	resultAnnotations := make([]string, 0, 3)
+	reflectedAnnotations := reflect.ValueOf(trace.FromContext(ctx)).Elem().Field(4).Elem().Field(0).Slice(0, 3)
+	for i := 0; i < 3; i++ {
+		resultAnnotations = append(resultAnnotations, fmt.Sprint(reflectedAnnotations.Index(i).Elem().Field(1)))
+	}
+
+	assert.Equal(t, expectedAnnotations, resultAnnotations)
+}
+
+func TestMetadataToAttributes(t *testing.T) {
+	expected := []trace.Attribute{trace.StringAttribute(fmt.Sprint("prefix.", expectedStr), "test value")}
+	result := metadataToAttributes(metadata.MD{expectedStr: {"test value"}}, "prefix.", defaultMetadataMatcher)
+	assert.Equal(t, expected, result)
+}
+
+func TestPayloadToAttributes(t *testing.T) {
+	expected := trace.StringAttribute(expectedStr, "\"test value\"")
+	result, ok, err := payloadToAttributes(expectedStr, "test value", 12)
+	assert.NoError(t, err)
+	assert.False(t, ok)
+	assert.Equal(t, expected, result[0])
+}
+
+func TestDefaultMetadataMatcher(t *testing.T) {
+	resultStr, ok := defaultMetadataMatcher(expectedStr)
+	assert.True(t, ok)
+	assert.Equal(t, expectedStr, resultStr)
+}
+
+func TestAlwaysGRPC(t *testing.T) {
+	assert.True(t, AlwaysGRPC(nil, nil))
+}

tracing/http.go

@@ -22,7 +22,7 @@ const (
 	ResponseHeaderAnnotationPrefix = "response.header."
 
 	//ResponseTrailerAnnotationPrefix is a prefix which is added to each response header attribute
-	ResponseTrailerAnnotationPrefix = "request.trailer."
+	ResponseTrailerAnnotationPrefix = "response.trailer."
 
 	//RequestPayloadAnnotationKey is a key under which request payload stored in span
 	RequestPayloadAnnotationKey = "request.payload"
@@ -94,7 +94,7 @@ func WithPayloadAnnotation(f func(*http.Request) bool) HTTPOption {
 	}
 }
 
-//WithHTTPPayloadSize limit payload size propogated to span
+//WithHTTPPayloadSize limit payload size propagated to span
 //in case payload exceeds limit, payload truncated and
 //annotation payload.truncated=true added into span
 func WithHTTPPayloadSize(maxSize int) HTTPOption {
@@ -125,7 +125,7 @@ func NewMiddleware(ops ...HTTPOption) func(http.Handler) http.Handler {
 //Check that &Handler comply with http.Handler interface
 var _ http.Handler = &Handler{}
 
-//Handler is a opencensus http plugin wrapper which do some usefull things to reach traces
+//Handler is a opencensus http plugin wrapper which do some useful things to reach traces
 type Handler struct {
 	child http.Handler
 
@@ -219,7 +219,7 @@ type responseBodyWrapper struct {
 }
 
 func (w *responseBodyWrapper) Write(b []byte) (int, error) {
-	//In case we recieve an error from Writing into buffer we just skip it
+	//In case we receive an error from Writing into buffer we just skip it
 	//because adding payload to span is not so critical as provide response
 	_, _ = w.buffer.Write(b)
 	return w.ResponseWriter.Write(b)