Added test

test/stats_test.go

@@ -0,0 +1,95 @@
+package test
+
+import (
+	"context"
+	"net"
+	"testing"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/interop"
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/stats"
+)
+
+// TestPeerForClientStatsHandler tests the scenario where stats handler
+// (having peer as part of its struct) has peer enriched as part of
+// stream context
+func (s) TestPeerForClientStatsHandler(t *testing.T) {
+	spy := &handlerSpy{}
+
+	// Start server.
+	l, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	grpcServer := grpc.NewServer()
+	testgrpc.RegisterTestServiceServer(grpcServer, interop.NewTestServer())
+	errCh := make(chan error)
+	go func() {
+		errCh <- grpcServer.Serve(l)
+	}()
+	t.Cleanup(func() {
+		grpcServer.Stop()
+		if err := <-errCh; err != nil {
+			t.Error(err)
+		}
+	})
+
+	// Create client with stats handler and do some calls.
+	conn, err := grpc.Dial(
+		l.Addr().String(),
+		grpc.WithBlock(),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithStatsHandler(spy))
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		if err := conn.Close(); err != nil {
+			t.Error(err)
+		}
+	})
+
+	ctx := context.Background()
+	client := testgrpc.NewTestServiceClient(conn)
+	interop.DoClientStreaming(ctx, client)
+
+	// Assert if peer is populated for each stats type.
+	for _, callbackArgs := range spy.Args {
+		if callbackArgs.Peer == nil {
+			switch callbackArgs.RPCStats.(type) {
+			case *stats.Begin:
+				continue
+			default:
+			}
+			t.Errorf("peer not populated for: %T", callbackArgs.RPCStats)
+		}
+	}
+}
+
+type peerStats struct {
+	RPCStats stats.RPCStats
+	Peer     *peer.Peer
+}
+
+type handlerSpy struct {
+	Args []peerStats
+}
+
+func (h *handlerSpy) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+	return ctx
+}
+
+func (h *handlerSpy) HandleRPC(ctx context.Context, rs stats.RPCStats) {
+	p, _ := peer.FromContext(ctx)
+	h.Args = append(h.Args, peerStats{rs, p})
+}
+
+func (h *handlerSpy) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
+	return ctx
+}
+
+func (h *handlerSpy) HandleConn(context.Context, stats.ConnStats) {
+}