go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/grpc/grpcmon/client_test.go (about)

     1  // Copyright 2016 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package grpcmon
    16  
    17  import (
    18  	"context"
    19  	"net"
    20  	"testing"
    21  	"time"
    22  
    23  	"google.golang.org/grpc"
    24  	"google.golang.org/grpc/codes"
    25  	"google.golang.org/grpc/credentials/insecure"
    26  	"google.golang.org/grpc/status"
    27  
    28  	"go.chromium.org/luci/common/tsmon/distribution"
    29  
    30  	. "github.com/smartystreets/goconvey/convey"
    31  )
    32  
    33  type echoService struct {
    34  	err error
    35  }
    36  
    37  func (s *echoService) Say(ctx context.Context, req *SayRequest) (*SayResponse, error) {
    38  	return &SayResponse{Msg: req.GetMsg()}, s.err
    39  }
    40  
    41  func TestClientRPCStatsMonitor(t *testing.T) {
    42  	method := "/grpcmon.Echo/Say"
    43  	fields := func(fs ...any) (ret []any) {
    44  		return append([]any{method}, fs...)
    45  	}
    46  
    47  	Convey("ClientRPCStatsMonitor", t, func() {
    48  		// spin up a server
    49  		srv, svc := grpc.NewServer(), &echoService{}
    50  		RegisterEchoServer(srv, svc)
    51  		l, err := net.Listen("tcp", "localhost:0")
    52  		So(err, ShouldBeNil)
    53  		go func() { _ = srv.Serve(l) }()
    54  		defer srv.Stop()
    55  
    56  		// construct a client
    57  		conn, err := grpc.Dial(
    58  			l.Addr().String(),
    59  			grpc.WithTransportCredentials(insecure.NewCredentials()),
    60  			grpc.WithBlock(),
    61  			grpc.WithStatsHandler(&ClientRPCStatsMonitor{}),
    62  		)
    63  		So(err, ShouldBeNil)
    64  		defer func() { So(conn.Close(), ShouldBeNil) }()
    65  		client := NewEchoClient(conn)
    66  		ctx, memStore := testContext()
    67  
    68  		run := func(err error, msg string) {
    69  			svc.err = err
    70  			resp, rerr := client.Say(ctx, &SayRequest{Msg: msg})
    71  			if err == nil {
    72  				So(rerr, ShouldBeNil)
    73  				So(resp.GetMsg(), ShouldEqual, msg)
    74  			} else {
    75  				So(rerr.Error(), ShouldEqual, err.Error())
    76  			}
    77  		}
    78  		Convey("Captures count and duration", func() {
    79  			count := func(code string) int64 {
    80  				val := memStore.Get(ctx, grpcClientCount, time.Time{}, fields(code))
    81  				So(val, ShouldNotBeNil)
    82  				return val.(int64)
    83  			}
    84  			duration := func(code string) any {
    85  				return memStore.Get(ctx, grpcClientDuration, time.Time{}, fields(code))
    86  			}
    87  
    88  			// grpc uses time.Now() to assign a value to
    89  			// grpc.End.{BeginTime, EndTime}, and we cannot stub it out.
    90  			//
    91  			// Therefore, this only checks the duration has been set or not.
    92  			// i.e., nil or not.
    93  			So(duration("OK"), ShouldBeNil)
    94  			run(nil, "echo!")
    95  			So(count("OK"), ShouldEqual, 1)
    96  			So(duration("OK"), ShouldNotBeNil)
    97  
    98  			So(duration("PERMISSION_DENIED"), ShouldBeNil)
    99  			run(status.Error(codes.PermissionDenied, "no permission"), "echo!")
   100  			So(count("PERMISSION_DENIED"), ShouldEqual, 1)
   101  			So(duration("PERMISSION_DENIED"), ShouldNotBeNil)
   102  
   103  			So(duration("UNAUTHENTICATED"), ShouldBeNil)
   104  			run(status.Error(codes.Unauthenticated, "no auth"), "echo!")
   105  			So(count("UNAUTHENTICATED"), ShouldEqual, 1)
   106  			So(duration("UNAUTHENTICATED"), ShouldNotBeNil)
   107  		})
   108  
   109  		Convey("Captures sent/received messages", func() {
   110  			count := func(code string) (float64, float64) {
   111  				sent := memStore.Get(ctx, grpcClientSentMsg, time.Time{}, fields())
   112  				So(sent, ShouldNotBeNil)
   113  				recv := memStore.Get(ctx, grpcClientRecvMsg, time.Time{}, fields())
   114  				So(recv, ShouldNotBeNil)
   115  				return sent.(*distribution.Distribution).Sum(), recv.(*distribution.Distribution).Sum()
   116  			}
   117  			bytes := func(code string) (float64, float64) {
   118  				sent := memStore.Get(ctx, grpcClientSentByte, time.Time{}, fields())
   119  				So(sent, ShouldNotBeNil)
   120  				recv := memStore.Get(ctx, grpcClientRecvByte, time.Time{}, fields())
   121  				So(recv, ShouldNotBeNil)
   122  				return sent.(*distribution.Distribution).Sum(), recv.(*distribution.Distribution).Sum()
   123  			}
   124  
   125  			run(nil, "echo!")
   126  			sentCount, recvCount := count("OK")
   127  			So(sentCount, ShouldEqual, 1)
   128  			So(recvCount, ShouldEqual, 1)
   129  
   130  			sentBytes, recvBytes := bytes("OK")
   131  			So(sentBytes, ShouldBeGreaterThan, 0)
   132  			So(recvBytes, ShouldBeGreaterThan, 0)
   133  		})
   134  	})
   135  }