From fde22f93f2b1221dabccdb766f0c32aec727ebc7 Mon Sep 17 00:00:00 2001
From: Abhishek Ranjan <aranjans@google.com>
Date: Fri, 19 Apr 2024 12:11:45 +0530
Subject: [PATCH] final fix

---
 test/stats_test.go | 59 ++++++++++++++++++++++++----------------------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/test/stats_test.go b/test/stats_test.go
index cefbf108e80d..f95ba3a5885a 100644
--- a/test/stats_test.go
+++ b/test/stats_test.go
@@ -32,91 +32,94 @@ import (
 	"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.
+// TestPeerForClientStatsHandler configures a stats handler that
+// verifies that peer is sent for OutPayload, InPayload, End
+// stats handlers.
 func (s) TestPeerForClientStatsHandler(t *testing.T) {
-	spy := &handlerSpy{}
+	statsHandler := &peerStatsHandler{}
 
 	// Start server.
 	l, err := net.Listen("tcp", "localhost:0")
 	if err != nil {
 		t.Fatal(err)
 	}
-	grpcServer := grpc.NewServer()
-	testgrpc.RegisterTestServiceServer(grpcServer, interop.NewTestServer())
+	s := grpc.NewServer()
+	testgrpc.RegisterTestServiceServer(s, interop.NewTestServer())
 	errCh := make(chan error)
 	go func() {
-		errCh <- grpcServer.Serve(l)
+		errCh <- s.Serve(l)
 	}()
-	t.Cleanup(func() {
-		grpcServer.Stop()
+	defer func() {
+		s.Stop()
 		if err := <-errCh; err != nil {
 			t.Error(err)
 		}
-	})
+	}()
 
 	// Create client with stats handler and do some calls.
-	conn, err := grpc.NewClient(
+	cc, err := grpc.NewClient(
 		l.Addr().String(),
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
-		grpc.WithStatsHandler(spy))
+		grpc.WithStatsHandler(statsHandler))
 	if err != nil {
 		t.Fatal(err)
 	}
 	t.Cleanup(func() {
-		if err := conn.Close(); err != nil {
+		if err := cc.Close(); err != nil {
 			t.Error(err)
 		}
 	})
 
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
-	client := testgrpc.NewTestServiceClient(conn)
+	client := testgrpc.NewTestServiceClient(cc)
 	interop.DoEmptyUnaryCall(ctx, client)
 	interop.DoLargeUnaryCall(ctx, client)
 	interop.DoClientStreaming(ctx, client)
 	interop.DoServerStreaming(ctx, client)
 	interop.DoPingPong(ctx, client)
 
-	// Assert if peer is populated for each stats type.
-	for _, callbackArgs := range spy.Args {
-		if callbackArgs.Peer == nil {
-			switch callbackArgs.RPCStats.(type) {
+	// Assert if peer is populated for each stats type except
+	// for Begin(as RPC begins even before we got the peer, and
+	// PickerUpdated(as PickerUpdated don't have any remote peer
+	// information since it happens without transport)
+	for _, callbackArgs := range statsHandler.Args {
+		if callbackArgs.peer == nil {
+			switch callbackArgs.rpcStats.(type) {
 			case *stats.Begin:
 				continue
 			case *stats.PickerUpdated:
 				continue
 			default:
 			}
-			t.Errorf("peer not populated for: %T", callbackArgs.RPCStats)
+			t.Errorf("peer not populated for: %T", callbackArgs.rpcStats)
 		}
 	}
 }
 
 type peerStats struct {
-	RPCStats stats.RPCStats
-	Peer     *peer.Peer
+	rpcStats stats.RPCStats
+	peer     *peer.Peer
 }
 
-type handlerSpy struct {
+type peerStatsHandler struct {
 	Args []peerStats
 	mu   sync.Mutex
 }
 
-func (h *handlerSpy) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+func (h *peerStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
 	return ctx
 }
 
-func (h *handlerSpy) HandleRPC(ctx context.Context, rs stats.RPCStats) {
-	h.mu.Lock()
+func (h *peerStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 	p, _ := peer.FromContext(ctx)
+	h.mu.Lock()
 	h.Args = append(h.Args, peerStats{rs, p})
-	h.mu.Unlock()
+	defer h.mu.Unlock()
 }
 
-func (h *handlerSpy) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
+func (h *peerStatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
 	return ctx
 }
 
-func (h *handlerSpy) HandleConn(context.Context, stats.ConnStats) {}
+func (h *peerStatsHandler) HandleConn(context.Context, stats.ConnStats) {}