github.com/blend/go-sdk@v1.20220411.3/tracing/grpctrace/main_test.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package grpctrace
     9  
    10  import (
    11  	"context"
    12  	"net"
    13  	"testing"
    14  	"time"
    15  
    16  	opentracing "github.com/opentracing/opentracing-go"
    17  	"github.com/opentracing/opentracing-go/mocktracer"
    18  	"google.golang.org/grpc"
    19  
    20  	"github.com/blend/go-sdk/assert"
    21  	"github.com/blend/go-sdk/grpcutil"
    22  	"github.com/blend/go-sdk/grpcutil/calculator"
    23  	"github.com/blend/go-sdk/tracing"
    24  
    25  	v1 "github.com/blend/go-sdk/grpcutil/calculator/v1"
    26  )
    27  
    28  func Test_Tracing_ServerUnary(t *testing.T) {
    29  	assert := assert.New(t)
    30  
    31  	mockTracer := mocktracer.New()
    32  	tracer := Tracer(mockTracer)
    33  
    34  	// start mocked server with tracing enabled
    35  	socketListener, err := net.Listen("tcp", "127.0.0.1:")
    36  	assert.Nil(err)
    37  	defer func() { _ = socketListener.Close() }()
    38  
    39  	server := grpc.NewServer(grpc.UnaryInterceptor(grpcutil.TracedServerUnary(tracer)))
    40  	v1.RegisterCalculatorServer(server, new(calculator.Server))
    41  	go func() { _ = server.Serve(socketListener) }()
    42  
    43  	conn, err := grpc.Dial(socketListener.Addr().String(), grpc.WithInsecure())
    44  	assert.Nil(err)
    45  	res, err := v1.NewCalculatorClient(conn).Add(context.Background(), &v1.Numbers{Values: []float64{1, 2, 3, 4}})
    46  	assert.Nil(err)
    47  	assert.Equal(10, res.Value)
    48  
    49  	assert.Len(mockTracer.FinishedSpans(), 1)
    50  	span := mockTracer.FinishedSpans()[0]
    51  	assert.Equal(tracing.SpanTypeRPC, span.Tags()["span.type"])
    52  	assert.Equal("grpc.server.unary", span.OperationName)
    53  	assert.Equal("/v1.Calculator/Add", span.Tags()[tracing.TagKeyResourceName])
    54  	assert.Equal("server", span.Tags()[tracing.TagKeyGRPCRole])
    55  	assert.Equal("unary", span.Tags()[tracing.TagKeyGRPCCallingConvention])
    56  }
    57  
    58  func Test_Tracing_ServerStream(t *testing.T) {
    59  	assert := assert.New(t)
    60  
    61  	mockTracer := mocktracer.New()
    62  	tracer := Tracer(mockTracer)
    63  
    64  	// start mocked server with tracing enabled
    65  	socketListener, err := net.Listen("tcp", "127.0.0.1:")
    66  	assert.Nil(err)
    67  	defer func() { _ = socketListener.Close() }()
    68  
    69  	server := grpc.NewServer(grpc.StreamInterceptor(grpcutil.TracedServerStream(tracer)))
    70  	v1.RegisterCalculatorServer(server, new(calculator.Server))
    71  	go func() { _ = server.Serve(socketListener) }()
    72  
    73  	conn, err := grpc.Dial(socketListener.Addr().String(), grpc.WithInsecure())
    74  	assert.Nil(err)
    75  	stream, err := v1.NewCalculatorClient(conn).AddStream(context.Background())
    76  	assert.Nil(err)
    77  
    78  	assert.Nil(stream.Send(&v1.Number{Value: 1}))
    79  	assert.Nil(stream.Send(&v1.Number{Value: 2}))
    80  	assert.Nil(stream.Send(&v1.Number{Value: 3}))
    81  	assert.Nil(stream.Send(&v1.Number{Value: 4}))
    82  
    83  	res, err := stream.CloseAndRecv()
    84  	assert.Nil(err)
    85  	assert.Equal(10, res.Value)
    86  
    87  	assert.Len(mockTracer.FinishedSpans(), 1)
    88  	span := mockTracer.FinishedSpans()[0]
    89  	assert.Equal(tracing.SpanTypeRPC, span.Tags()["span.type"])
    90  	assert.Equal("grpc.server.stream", span.OperationName)
    91  	assert.Equal("/v1.Calculator/AddStream", span.Tags()[tracing.TagKeyResourceName])
    92  	assert.Equal("server", span.Tags()[tracing.TagKeyGRPCRole])
    93  	assert.Equal("stream", span.Tags()[tracing.TagKeyGRPCCallingConvention])
    94  }
    95  
    96  func Test_Tracing_ClientServerUnary(t *testing.T) {
    97  	assert := assert.New(t)
    98  
    99  	serverTraceCollector := mocktracer.New()
   100  	serverTracer := Tracer(serverTraceCollector)
   101  
   102  	clientTraceCollector := mocktracer.New()
   103  	clientTracer := Tracer(clientTraceCollector)
   104  
   105  	// start mocked server with tracing enabled
   106  	socketListener, err := net.Listen("tcp", "127.0.0.1:")
   107  	assert.Nil(err)
   108  	defer func() { _ = socketListener.Close() }()
   109  
   110  	server := grpc.NewServer(grpc.UnaryInterceptor(grpcutil.TracedServerUnary(serverTracer)))
   111  	v1.RegisterCalculatorServer(server, new(calculator.Server))
   112  	go func() { _ = server.Serve(socketListener) }()
   113  
   114  	conn, err := grpc.Dial(socketListener.Addr().String(), grpc.WithInsecure(), grpc.WithUnaryInterceptor(grpcutil.TracedClientUnary(clientTracer)))
   115  	assert.Nil(err)
   116  	res, err := v1.NewCalculatorClient(conn).Add(context.Background(), &v1.Numbers{Values: []float64{1, 2, 3, 4}})
   117  	assert.Nil(err)
   118  	assert.Equal(10, res.Value)
   119  
   120  	assert.Len(serverTraceCollector.FinishedSpans(), 1)
   121  	assert.Len(clientTraceCollector.FinishedSpans(), 1)
   122  
   123  	// server
   124  	serverSpan := serverTraceCollector.FinishedSpans()[0]
   125  	assert.NotZero(serverSpan.ParentID)
   126  	assert.Equal(tracing.SpanTypeRPC, serverSpan.Tags()["span.type"])
   127  	assert.Equal("grpc.server.unary", serverSpan.OperationName)
   128  	assert.Equal("/v1.Calculator/Add", serverSpan.Tags()[tracing.TagKeyResourceName])
   129  	assert.Equal("server", serverSpan.Tags()[tracing.TagKeyGRPCRole])
   130  	assert.Equal("unary", serverSpan.Tags()[tracing.TagKeyGRPCCallingConvention])
   131  
   132  	// client
   133  	clientSpan := clientTraceCollector.FinishedSpans()[0]
   134  	assert.Zero(clientSpan.ParentID)
   135  	assert.Equal(tracing.SpanTypeRPC, clientSpan.Tags()["span.type"])
   136  	assert.Equal("grpc.client.unary", clientSpan.OperationName)
   137  	assert.Equal("/v1.Calculator/Add", clientSpan.Tags()[tracing.TagKeyResourceName])
   138  	assert.Equal("client", clientSpan.Tags()[tracing.TagKeyGRPCRole])
   139  	assert.Equal("unary", clientSpan.Tags()[tracing.TagKeyGRPCCallingConvention])
   140  }
   141  
   142  func Test_Tracing_ClientServerStream(t *testing.T) {
   143  	assert := assert.New(t)
   144  
   145  	mockTracer := mocktracer.New()
   146  	tracer := Tracer(mockTracer)
   147  
   148  	// start mocked server with tracing enabled
   149  	socketListener, err := net.Listen("tcp", "127.0.0.1:")
   150  	assert.Nil(err)
   151  	defer func() { _ = socketListener.Close() }()
   152  
   153  	server := grpc.NewServer(grpc.StreamInterceptor(grpcutil.TracedServerStream(tracer)))
   154  	v1.RegisterCalculatorServer(server, new(calculator.Server))
   155  	go func() { _ = server.Serve(socketListener) }()
   156  
   157  	conn, err := grpc.Dial(socketListener.Addr().String(), grpc.WithInsecure(), grpc.WithStreamInterceptor(grpcutil.TracedClientStream(tracer)))
   158  	assert.Nil(err)
   159  	stream, err := v1.NewCalculatorClient(conn).AddStream(context.Background())
   160  	assert.Nil(err)
   161  
   162  	assert.Nil(stream.Send(&v1.Number{Value: 1}))
   163  	assert.Nil(stream.Send(&v1.Number{Value: 2}))
   164  	assert.Nil(stream.Send(&v1.Number{Value: 3}))
   165  	assert.Nil(stream.Send(&v1.Number{Value: 4}))
   166  
   167  	res, err := stream.CloseAndRecv()
   168  	assert.Nil(err)
   169  	assert.Equal(10, res.Value)
   170  
   171  	assert.Len(mockTracer.FinishedSpans(), 2)
   172  	assert.Equal(tracing.SpanTypeRPC, mockTracer.FinishedSpans()[0].Tags()["span.type"])
   173  	assert.Equal("grpc.client.stream", mockTracer.FinishedSpans()[0].OperationName)
   174  	assert.Equal("/v1.Calculator/AddStream", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyResourceName])
   175  	assert.Equal("client", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyGRPCRole])
   176  	assert.Equal("stream", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyGRPCCallingConvention])
   177  
   178  	assert.NotZero(mockTracer.FinishedSpans()[1].ParentID)
   179  	assert.Equal(tracing.SpanTypeRPC, mockTracer.FinishedSpans()[1].Tags()["span.type"])
   180  	assert.Equal("grpc.server.stream", mockTracer.FinishedSpans()[1].OperationName)
   181  	assert.Equal("/v1.Calculator/AddStream", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyResourceName])
   182  	assert.Equal("server", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyGRPCRole])
   183  	assert.Equal("stream", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyGRPCCallingConvention])
   184  }
   185  
   186  func Test_Tracing_ParentClientServerUnary(t *testing.T) {
   187  	assert := assert.New(t)
   188  
   189  	mockTracer := mocktracer.New()
   190  	tracer := Tracer(mockTracer)
   191  
   192  	// start mocked server with tracing enabled
   193  	socketListener, err := net.Listen("tcp", "127.0.0.1:")
   194  	assert.Nil(err)
   195  	defer func() { _ = socketListener.Close() }()
   196  
   197  	server := grpc.NewServer(grpc.UnaryInterceptor(grpcutil.TracedServerUnary(tracer)))
   198  	v1.RegisterCalculatorServer(server, new(calculator.Server))
   199  	go func() { _ = server.Serve(socketListener) }()
   200  
   201  	outerSpan, ctx := tracing.StartSpanFromContext(context.Background(), mockTracer, tracing.OperationHTTPRequest,
   202  		opentracing.Tag{Key: tracing.TagKeyResourceName, Value: "/foo"},
   203  		opentracing.Tag{Key: tracing.TagKeySpanType, Value: tracing.SpanTypeWeb},
   204  		opentracing.StartTime(time.Now().UTC()),
   205  	)
   206  
   207  	conn, err := grpc.Dial(socketListener.Addr().String(), grpc.WithInsecure(), grpc.WithUnaryInterceptor(grpcutil.TracedClientUnary(tracer)))
   208  	assert.Nil(err)
   209  	res, err := v1.NewCalculatorClient(conn).Add(ctx, &v1.Numbers{Values: []float64{1, 2, 3, 4}})
   210  	assert.Nil(err)
   211  	assert.Equal(10, res.Value)
   212  
   213  	// finish the outer span ...
   214  	outerSpan.Finish()
   215  
   216  	assert.Len(mockTracer.FinishedSpans(), 3)
   217  
   218  	// server
   219  	assert.NotZero(mockTracer.FinishedSpans()[0].ParentID)
   220  	assert.Equal(tracing.SpanTypeRPC, mockTracer.FinishedSpans()[0].Tags()["span.type"])
   221  	assert.Equal("grpc.server.unary", mockTracer.FinishedSpans()[0].OperationName)
   222  	assert.Equal("/v1.Calculator/Add", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyResourceName])
   223  	assert.Equal("server", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyGRPCRole])
   224  	assert.Equal("unary", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyGRPCCallingConvention])
   225  
   226  	// client
   227  	assert.NotZero(mockTracer.FinishedSpans()[1].ParentID)
   228  	assert.Equal(tracing.SpanTypeRPC, mockTracer.FinishedSpans()[1].Tags()["span.type"])
   229  	assert.Equal("grpc.client.unary", mockTracer.FinishedSpans()[1].OperationName)
   230  	assert.Equal("/v1.Calculator/Add", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyResourceName])
   231  	assert.Equal("client", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyGRPCRole])
   232  	assert.Equal("unary", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyGRPCCallingConvention])
   233  }
   234  
   235  func Test_Tracing_ParentClientServerStream(t *testing.T) {
   236  	assert := assert.New(t)
   237  
   238  	mockTracer := mocktracer.New()
   239  	tracer := Tracer(mockTracer)
   240  
   241  	// start mocked server with tracing enabled
   242  	socketListener, err := net.Listen("tcp", "127.0.0.1:")
   243  	assert.Nil(err)
   244  	defer func() { _ = socketListener.Close() }()
   245  
   246  	server := grpc.NewServer(grpc.StreamInterceptor(grpcutil.TracedServerStream(tracer)))
   247  	v1.RegisterCalculatorServer(server, new(calculator.Server))
   248  	go func() { _ = server.Serve(socketListener) }()
   249  
   250  	conn, err := grpc.Dial(socketListener.Addr().String(), grpc.WithInsecure(), grpc.WithStreamInterceptor(grpcutil.TracedClientStream(tracer)))
   251  	assert.Nil(err)
   252  
   253  	outerSpan, ctx := tracing.StartSpanFromContext(context.Background(), mockTracer, tracing.OperationHTTPRequest,
   254  		opentracing.Tag{Key: tracing.TagKeyResourceName, Value: "/foo"},
   255  		opentracing.Tag{Key: tracing.TagKeySpanType, Value: tracing.SpanTypeWeb},
   256  		opentracing.StartTime(time.Now().UTC()),
   257  	)
   258  
   259  	stream, err := v1.NewCalculatorClient(conn).AddStream(ctx)
   260  	assert.Nil(err)
   261  
   262  	assert.Nil(stream.Send(&v1.Number{Value: 1}))
   263  	assert.Nil(stream.Send(&v1.Number{Value: 2}))
   264  	assert.Nil(stream.Send(&v1.Number{Value: 3}))
   265  	assert.Nil(stream.Send(&v1.Number{Value: 4}))
   266  
   267  	res, err := stream.CloseAndRecv()
   268  	assert.Nil(err)
   269  	assert.Equal(10, res.Value)
   270  
   271  	// finish the outer span ...
   272  	outerSpan.Finish()
   273  
   274  	assert.Len(mockTracer.FinishedSpans(), 3)
   275  
   276  	assert.NotZero(mockTracer.FinishedSpans()[0].ParentID)
   277  	assert.Equal(tracing.SpanTypeRPC, mockTracer.FinishedSpans()[0].Tags()["span.type"])
   278  	assert.Equal("grpc.client.stream", mockTracer.FinishedSpans()[0].OperationName)
   279  	assert.Equal("/v1.Calculator/AddStream", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyResourceName])
   280  	assert.Equal("client", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyGRPCRole])
   281  	assert.Equal("stream", mockTracer.FinishedSpans()[0].Tags()[tracing.TagKeyGRPCCallingConvention])
   282  
   283  	assert.NotZero(mockTracer.FinishedSpans()[1].ParentID)
   284  	assert.Equal(tracing.SpanTypeRPC, mockTracer.FinishedSpans()[1].Tags()["span.type"])
   285  	assert.Equal("grpc.server.stream", mockTracer.FinishedSpans()[1].OperationName)
   286  	assert.Equal("/v1.Calculator/AddStream", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyResourceName])
   287  	assert.Equal("server", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyGRPCRole])
   288  	assert.Equal("stream", mockTracer.FinishedSpans()[1].Tags()[tracing.TagKeyGRPCCallingConvention])
   289  }