File: grpc_test.go

package info (click to toggle)
golang-github-go-kit-kit 0.13.0-8
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 1,784 kB
  • sloc: sh: 22; makefile: 11
file content (111 lines) | stat: -rw-r--r-- 2,902 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
package zipkin_test

import (
	"context"
	"testing"

	zipkin "github.com/openzipkin/zipkin-go"
	"github.com/openzipkin/zipkin-go/propagation/b3"
	"github.com/openzipkin/zipkin-go/reporter/recorder"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	"github.com/go-kit/kit/endpoint"
	kitzipkin "github.com/go-kit/kit/tracing/zipkin"
	grpctransport "github.com/go-kit/kit/transport/grpc"
)

type dummy struct{}

func unaryInterceptor(
	ctx context.Context, method string, req, reply interface{},
	cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption,
) error {
	return nil
}

func TestGRPCClientTrace(t *testing.T) {
	rec := recorder.NewReporter()
	defer rec.Close()

	tr, _ := zipkin.NewTracer(rec)

	clientTracer := kitzipkin.GRPCClientTrace(tr)

	cc, err := grpc.Dial(
		"",
		grpc.WithUnaryInterceptor(unaryInterceptor),
		grpc.WithInsecure(),
	)
	if err != nil {
		t.Fatalf("unable to create gRPC dialer: %s", err.Error())
	}

	ep := grpctransport.NewClient(
		cc,
		"dummyService",
		"dummyMethod",
		func(context.Context, interface{}) (interface{}, error) { return nil, nil },
		func(context.Context, interface{}) (interface{}, error) { return nil, nil },
		dummy{},
		clientTracer,
	).Endpoint()

	parentSpan := tr.StartSpan("test")
	ctx := zipkin.NewContext(context.Background(), parentSpan)

	if _, err = ep(ctx, nil); err != nil {
		t.Errorf("unexpected error: %s", err.Error())
	}

	spans := rec.Flush()
	if want, have := 1, len(spans); want != have {
		t.Fatalf("incorrect number of spans, want %d, have %d", want, have)
	}

	if spans[0].SpanContext.ParentID == nil {
		t.Fatalf("incorrect parent ID, want %s have nil", parentSpan.Context().ID)
	}

	if want, have := parentSpan.Context().ID, *spans[0].SpanContext.ParentID; want != have {
		t.Fatalf("incorrect parent ID, want %s, have %s", want, have)
	}
}

func TestGRPCServerTrace(t *testing.T) {
	rec := recorder.NewReporter()
	defer rec.Close()

	tr, _ := zipkin.NewTracer(rec)

	serverTracer := kitzipkin.GRPCServerTrace(tr)

	server := grpctransport.NewServer(
		endpoint.Nop,
		func(context.Context, interface{}) (interface{}, error) { return nil, nil },
		func(context.Context, interface{}) (interface{}, error) { return nil, nil },
		serverTracer,
	)

	md := metadata.MD{}
	parentSpan := tr.StartSpan("test")

	b3.InjectGRPC(&md)(parentSpan.Context())

	ctx := metadata.NewIncomingContext(context.Background(), md)
	server.ServeGRPC(ctx, nil)

	spans := rec.Flush()

	if want, have := 1, len(spans); want != have {
		t.Fatalf("incorrect number of spans, want %d, have %d", want, have)
	}

	if want, have := parentSpan.Context().TraceID, spans[0].SpanContext.TraceID; want != have {
		t.Errorf("incorrect TraceID, want %+v, have %+v", want, have)
	}

	if want, have := parentSpan.Context().ID, spans[0].SpanContext.ID; want != have {
		t.Errorf("incorrect span ID, want %d, have %d", want, have)
	}
}