...

Source file src/github.com/go-kit/kit/tracing/opencensus/grpc_test.go

Documentation: github.com/go-kit/kit/tracing/opencensus

     1  package opencensus_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  
     8  	"go.opencensus.io/trace"
     9  	"go.opencensus.io/trace/propagation"
    10  	"google.golang.org/grpc"
    11  	"google.golang.org/grpc/codes"
    12  	"google.golang.org/grpc/metadata"
    13  
    14  	"github.com/go-kit/kit/endpoint"
    15  	ockit "github.com/go-kit/kit/tracing/opencensus"
    16  	grpctransport "github.com/go-kit/kit/transport/grpc"
    17  )
    18  
    19  type dummy struct{}
    20  
    21  const traceContextKey = "grpc-trace-bin"
    22  
    23  func unaryInterceptor(
    24  	ctx context.Context, method string, req, reply interface{},
    25  	cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption,
    26  ) error {
    27  	return nil
    28  }
    29  
    30  func TestGRPCClientTrace(t *testing.T) {
    31  	rec := &recordingExporter{}
    32  
    33  	trace.RegisterExporter(rec)
    34  
    35  	cc, err := grpc.Dial(
    36  		"",
    37  		grpc.WithUnaryInterceptor(unaryInterceptor),
    38  		grpc.WithInsecure(),
    39  	)
    40  	if err != nil {
    41  		t.Fatalf("unable to create gRPC dialer: %s", err.Error())
    42  	}
    43  
    44  	traces := []struct {
    45  		name string
    46  		err  error
    47  	}{
    48  		{"", nil},
    49  		{"CustomName", nil},
    50  		{"", errors.New("dummy-error")},
    51  	}
    52  
    53  	for _, tr := range traces {
    54  		clientTracer := ockit.GRPCClientTrace(
    55  			ockit.WithName(tr.name),
    56  			ockit.WithSampler(trace.AlwaysSample()),
    57  		)
    58  
    59  		ep := grpctransport.NewClient(
    60  			cc,
    61  			"dummyService",
    62  			"dummyMethod",
    63  			func(context.Context, interface{}) (interface{}, error) {
    64  				return nil, nil
    65  			},
    66  			func(context.Context, interface{}) (interface{}, error) {
    67  				return nil, tr.err
    68  			},
    69  			dummy{},
    70  			clientTracer,
    71  		).Endpoint()
    72  
    73  		ctx, parentSpan := trace.StartSpan(context.Background(), "test")
    74  
    75  		_, err = ep(ctx, nil)
    76  		if want, have := tr.err, err; want != have {
    77  			t.Fatalf("unexpected error, want %s, have %s", tr.err.Error(), err.Error())
    78  		}
    79  
    80  		spans := rec.Flush()
    81  		if want, have := 1, len(spans); want != have {
    82  			t.Fatalf("incorrect number of spans, want %d, have %d", want, have)
    83  		}
    84  		span := spans[0]
    85  		if want, have := parentSpan.SpanContext().SpanID, span.ParentSpanID; want != have {
    86  			t.Errorf("incorrect parent ID, want %s, have %s", want, have)
    87  		}
    88  
    89  		if want, have := tr.name, span.Name; want != have && want != "" {
    90  			t.Errorf("incorrect span name, want %s, have %s", want, have)
    91  		}
    92  
    93  		if want, have := "/dummyService/dummyMethod", span.Name; want != have && tr.name == "" {
    94  			t.Errorf("incorrect span name, want %s, have %s", want, have)
    95  		}
    96  
    97  		code := trace.StatusCodeOK
    98  		if tr.err != nil {
    99  			code = trace.StatusCodeUnknown
   100  
   101  			if want, have := err.Error(), span.Status.Message; want != have {
   102  				t.Errorf("incorrect span status msg, want %s, have %s", want, have)
   103  			}
   104  		}
   105  
   106  		if want, have := int32(code), span.Status.Code; want != have {
   107  			t.Errorf("incorrect span status code, want %d, have %d", want, have)
   108  		}
   109  	}
   110  }
   111  
   112  func TestGRPCServerTrace(t *testing.T) {
   113  	rec := &recordingExporter{}
   114  
   115  	trace.RegisterExporter(rec)
   116  
   117  	traces := []struct {
   118  		useParent bool
   119  		name      string
   120  		err       error
   121  	}{
   122  		{false, "", nil},
   123  		{true, "", nil},
   124  		{true, "CustomName", nil},
   125  		{true, "", errors.New("dummy-error")},
   126  	}
   127  
   128  	for _, tr := range traces {
   129  		var (
   130  			ctx        = context.Background()
   131  			parentSpan *trace.Span
   132  		)
   133  
   134  		server := grpctransport.NewServer(
   135  			endpoint.Nop,
   136  			func(context.Context, interface{}) (interface{}, error) {
   137  				return nil, nil
   138  			},
   139  			func(context.Context, interface{}) (interface{}, error) {
   140  				return nil, tr.err
   141  			},
   142  			ockit.GRPCServerTrace(
   143  				ockit.WithName(tr.name),
   144  				ockit.WithSampler(trace.AlwaysSample()),
   145  			),
   146  		)
   147  
   148  		if tr.useParent {
   149  			_, parentSpan = trace.StartSpan(context.Background(), "test")
   150  			traceContextBinary := propagation.Binary(parentSpan.SpanContext())
   151  
   152  			md := metadata.MD{}
   153  			md.Set(traceContextKey, string(traceContextBinary))
   154  			ctx = metadata.NewIncomingContext(ctx, md)
   155  		}
   156  
   157  		server.ServeGRPC(ctx, nil)
   158  
   159  		spans := rec.Flush()
   160  
   161  		if want, have := 1, len(spans); want != have {
   162  			t.Fatalf("incorrect number of spans, want %d, have %d", want, have)
   163  		}
   164  
   165  		if tr.useParent {
   166  			if want, have := parentSpan.SpanContext().TraceID, spans[0].TraceID; want != have {
   167  				t.Errorf("incorrect trace ID, want %s, have %s", want, have)
   168  			}
   169  
   170  			if want, have := parentSpan.SpanContext().SpanID, spans[0].ParentSpanID; want != have {
   171  				t.Errorf("incorrect span ID, want %s, have %s", want, have)
   172  			}
   173  		}
   174  
   175  		if want, have := tr.name, spans[0].Name; want != have && want != "" {
   176  			t.Errorf("incorrect span name, want %s, have %s", want, have)
   177  		}
   178  
   179  		if tr.err != nil {
   180  			if want, have := int32(codes.Internal), spans[0].Status.Code; want != have {
   181  				t.Errorf("incorrect span status code, want %d, have %d", want, have)
   182  			}
   183  
   184  			if want, have := tr.err.Error(), spans[0].Status.Message; want != have {
   185  				t.Errorf("incorrect span status message, want %s, have %s", want, have)
   186  			}
   187  		}
   188  	}
   189  }
   190  

View as plain text