github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/grpc/interop/test_utils.go (about)

     1  /*
     2   *
     3   * Copyright 2014 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  // Package interop contains functions used by interop client/server.
    20  package interop
    21  
    22  import (
    23  	"bytes"
    24  	"context"
    25  	"fmt"
    26  	"io"
    27  	"io/ioutil"
    28  	"os"
    29  	"strings"
    30  	"time"
    31  
    32  	"github.com/golang/protobuf/proto"
    33  	grpc "github.com/hxx258456/ccgo/grpc"
    34  	"github.com/hxx258456/ccgo/grpc/benchmark/stats"
    35  	"github.com/hxx258456/ccgo/grpc/codes"
    36  	"github.com/hxx258456/ccgo/grpc/grpclog"
    37  	"github.com/hxx258456/ccgo/grpc/metadata"
    38  	"github.com/hxx258456/ccgo/grpc/status"
    39  	"golang.org/x/oauth2"
    40  	"golang.org/x/oauth2/google"
    41  
    42  	testgrpc "github.com/hxx258456/ccgo/grpc/interop/grpc_testing"
    43  	testpb "github.com/hxx258456/ccgo/grpc/interop/grpc_testing"
    44  )
    45  
    46  var (
    47  	reqSizes            = []int{27182, 8, 1828, 45904}
    48  	respSizes           = []int{31415, 9, 2653, 58979}
    49  	largeReqSize        = 271828
    50  	largeRespSize       = 314159
    51  	initialMetadataKey  = "x-grpc-test-echo-initial"
    52  	trailingMetadataKey = "x-grpc-test-echo-trailing-bin"
    53  
    54  	logger = grpclog.Component("interop")
    55  )
    56  
    57  // ClientNewPayload returns a payload of the given type and size.
    58  func ClientNewPayload(t testpb.PayloadType, size int) *testpb.Payload {
    59  	if size < 0 {
    60  		logger.Fatalf("Requested a response with invalid length %d", size)
    61  	}
    62  	body := make([]byte, size)
    63  	switch t {
    64  	case testpb.PayloadType_COMPRESSABLE:
    65  	default:
    66  		logger.Fatalf("Unsupported payload type: %d", t)
    67  	}
    68  	return &testpb.Payload{
    69  		Type: t,
    70  		Body: body,
    71  	}
    72  }
    73  
    74  // DoEmptyUnaryCall performs a unary RPC with empty request and response messages.
    75  func DoEmptyUnaryCall(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
    76  	reply, err := tc.EmptyCall(context.Background(), &testpb.Empty{}, args...)
    77  	if err != nil {
    78  		logger.Fatal("/TestService/EmptyCall RPC failed: ", err)
    79  	}
    80  	if !proto.Equal(&testpb.Empty{}, reply) {
    81  		logger.Fatalf("/TestService/EmptyCall receives %v, want %v", reply, testpb.Empty{})
    82  	}
    83  }
    84  
    85  // DoLargeUnaryCall performs a unary RPC with large payload in the request and response.
    86  func DoLargeUnaryCall(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
    87  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
    88  	req := &testpb.SimpleRequest{
    89  		ResponseType: testpb.PayloadType_COMPRESSABLE,
    90  		ResponseSize: int32(largeRespSize),
    91  		Payload:      pl,
    92  	}
    93  	reply, err := tc.UnaryCall(context.Background(), req, args...)
    94  	if err != nil {
    95  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
    96  	}
    97  	t := reply.GetPayload().GetType()
    98  	s := len(reply.GetPayload().GetBody())
    99  	if t != testpb.PayloadType_COMPRESSABLE || s != largeRespSize {
   100  		logger.Fatalf("Got the reply with type %d len %d; want %d, %d", t, s, testpb.PayloadType_COMPRESSABLE, largeRespSize)
   101  	}
   102  }
   103  
   104  // DoClientStreaming performs a client streaming RPC.
   105  func DoClientStreaming(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   106  	stream, err := tc.StreamingInputCall(context.Background(), args...)
   107  	if err != nil {
   108  		logger.Fatalf("%v.StreamingInputCall(_) = _, %v", tc, err)
   109  	}
   110  	var sum int
   111  	for _, s := range reqSizes {
   112  		pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, s)
   113  		req := &testpb.StreamingInputCallRequest{
   114  			Payload: pl,
   115  		}
   116  		if err := stream.Send(req); err != nil {
   117  			logger.Fatalf("%v has error %v while sending %v", stream, err, req)
   118  		}
   119  		sum += s
   120  	}
   121  	reply, err := stream.CloseAndRecv()
   122  	if err != nil {
   123  		logger.Fatalf("%v.CloseAndRecv() got error %v, want %v", stream, err, nil)
   124  	}
   125  	if reply.GetAggregatedPayloadSize() != int32(sum) {
   126  		logger.Fatalf("%v.CloseAndRecv().GetAggregatePayloadSize() = %v; want %v", stream, reply.GetAggregatedPayloadSize(), sum)
   127  	}
   128  }
   129  
   130  // DoServerStreaming performs a server streaming RPC.
   131  func DoServerStreaming(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   132  	respParam := make([]*testpb.ResponseParameters, len(respSizes))
   133  	for i, s := range respSizes {
   134  		respParam[i] = &testpb.ResponseParameters{
   135  			Size: int32(s),
   136  		}
   137  	}
   138  	req := &testpb.StreamingOutputCallRequest{
   139  		ResponseType:       testpb.PayloadType_COMPRESSABLE,
   140  		ResponseParameters: respParam,
   141  	}
   142  	stream, err := tc.StreamingOutputCall(context.Background(), req, args...)
   143  	if err != nil {
   144  		logger.Fatalf("%v.StreamingOutputCall(_) = _, %v", tc, err)
   145  	}
   146  	var rpcStatus error
   147  	var respCnt int
   148  	var index int
   149  	for {
   150  		reply, err := stream.Recv()
   151  		if err != nil {
   152  			rpcStatus = err
   153  			break
   154  		}
   155  		t := reply.GetPayload().GetType()
   156  		if t != testpb.PayloadType_COMPRESSABLE {
   157  			logger.Fatalf("Got the reply of type %d, want %d", t, testpb.PayloadType_COMPRESSABLE)
   158  		}
   159  		size := len(reply.GetPayload().GetBody())
   160  		if size != respSizes[index] {
   161  			logger.Fatalf("Got reply body of length %d, want %d", size, respSizes[index])
   162  		}
   163  		index++
   164  		respCnt++
   165  	}
   166  	if rpcStatus != io.EOF {
   167  		logger.Fatalf("Failed to finish the server streaming rpc: %v", rpcStatus)
   168  	}
   169  	if respCnt != len(respSizes) {
   170  		logger.Fatalf("Got %d reply, want %d", len(respSizes), respCnt)
   171  	}
   172  }
   173  
   174  // DoPingPong performs ping-pong style bi-directional streaming RPC.
   175  func DoPingPong(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   176  	stream, err := tc.FullDuplexCall(context.Background(), args...)
   177  	if err != nil {
   178  		logger.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
   179  	}
   180  	var index int
   181  	for index < len(reqSizes) {
   182  		respParam := []*testpb.ResponseParameters{
   183  			{
   184  				Size: int32(respSizes[index]),
   185  			},
   186  		}
   187  		pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, reqSizes[index])
   188  		req := &testpb.StreamingOutputCallRequest{
   189  			ResponseType:       testpb.PayloadType_COMPRESSABLE,
   190  			ResponseParameters: respParam,
   191  			Payload:            pl,
   192  		}
   193  		if err := stream.Send(req); err != nil {
   194  			logger.Fatalf("%v has error %v while sending %v", stream, err, req)
   195  		}
   196  		reply, err := stream.Recv()
   197  		if err != nil {
   198  			logger.Fatalf("%v.Recv() = %v", stream, err)
   199  		}
   200  		t := reply.GetPayload().GetType()
   201  		if t != testpb.PayloadType_COMPRESSABLE {
   202  			logger.Fatalf("Got the reply of type %d, want %d", t, testpb.PayloadType_COMPRESSABLE)
   203  		}
   204  		size := len(reply.GetPayload().GetBody())
   205  		if size != respSizes[index] {
   206  			logger.Fatalf("Got reply body of length %d, want %d", size, respSizes[index])
   207  		}
   208  		index++
   209  	}
   210  	if err := stream.CloseSend(); err != nil {
   211  		logger.Fatalf("%v.CloseSend() got %v, want %v", stream, err, nil)
   212  	}
   213  	if _, err := stream.Recv(); err != io.EOF {
   214  		logger.Fatalf("%v failed to complele the ping pong test: %v", stream, err)
   215  	}
   216  }
   217  
   218  // DoEmptyStream sets up a bi-directional streaming with zero message.
   219  func DoEmptyStream(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   220  	stream, err := tc.FullDuplexCall(context.Background(), args...)
   221  	if err != nil {
   222  		logger.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
   223  	}
   224  	if err := stream.CloseSend(); err != nil {
   225  		logger.Fatalf("%v.CloseSend() got %v, want %v", stream, err, nil)
   226  	}
   227  	if _, err := stream.Recv(); err != io.EOF {
   228  		logger.Fatalf("%v failed to complete the empty stream test: %v", stream, err)
   229  	}
   230  }
   231  
   232  // DoTimeoutOnSleepingServer performs an RPC on a sleep server which causes RPC timeout.
   233  func DoTimeoutOnSleepingServer(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   234  	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
   235  	defer cancel()
   236  	stream, err := tc.FullDuplexCall(ctx, args...)
   237  	if err != nil {
   238  		if status.Code(err) == codes.DeadlineExceeded {
   239  			return
   240  		}
   241  		logger.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
   242  	}
   243  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, 27182)
   244  	req := &testpb.StreamingOutputCallRequest{
   245  		ResponseType: testpb.PayloadType_COMPRESSABLE,
   246  		Payload:      pl,
   247  	}
   248  	if err := stream.Send(req); err != nil && err != io.EOF {
   249  		logger.Fatalf("%v.Send(_) = %v", stream, err)
   250  	}
   251  	if _, err := stream.Recv(); status.Code(err) != codes.DeadlineExceeded {
   252  		logger.Fatalf("%v.Recv() = _, %v, want error code %d", stream, err, codes.DeadlineExceeded)
   253  	}
   254  }
   255  
   256  // DoComputeEngineCreds performs a unary RPC with compute engine auth.
   257  func DoComputeEngineCreds(tc testgrpc.TestServiceClient, serviceAccount, oauthScope string) {
   258  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   259  	req := &testpb.SimpleRequest{
   260  		ResponseType:   testpb.PayloadType_COMPRESSABLE,
   261  		ResponseSize:   int32(largeRespSize),
   262  		Payload:        pl,
   263  		FillUsername:   true,
   264  		FillOauthScope: true,
   265  	}
   266  	reply, err := tc.UnaryCall(context.Background(), req)
   267  	if err != nil {
   268  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   269  	}
   270  	user := reply.GetUsername()
   271  	scope := reply.GetOauthScope()
   272  	if user != serviceAccount {
   273  		logger.Fatalf("Got user name %q, want %q.", user, serviceAccount)
   274  	}
   275  	if !strings.Contains(oauthScope, scope) {
   276  		logger.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, oauthScope)
   277  	}
   278  }
   279  
   280  func getServiceAccountJSONKey(keyFile string) []byte {
   281  	jsonKey, err := ioutil.ReadFile(keyFile)
   282  	if err != nil {
   283  		logger.Fatalf("Failed to read the service account key file: %v", err)
   284  	}
   285  	return jsonKey
   286  }
   287  
   288  // DoServiceAccountCreds performs a unary RPC with service account auth.
   289  func DoServiceAccountCreds(tc testgrpc.TestServiceClient, serviceAccountKeyFile, oauthScope string) {
   290  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   291  	req := &testpb.SimpleRequest{
   292  		ResponseType:   testpb.PayloadType_COMPRESSABLE,
   293  		ResponseSize:   int32(largeRespSize),
   294  		Payload:        pl,
   295  		FillUsername:   true,
   296  		FillOauthScope: true,
   297  	}
   298  	reply, err := tc.UnaryCall(context.Background(), req)
   299  	if err != nil {
   300  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   301  	}
   302  	jsonKey := getServiceAccountJSONKey(serviceAccountKeyFile)
   303  	user := reply.GetUsername()
   304  	scope := reply.GetOauthScope()
   305  	if !strings.Contains(string(jsonKey), user) {
   306  		logger.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
   307  	}
   308  	if !strings.Contains(oauthScope, scope) {
   309  		logger.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, oauthScope)
   310  	}
   311  }
   312  
   313  // DoJWTTokenCreds performs a unary RPC with JWT token auth.
   314  func DoJWTTokenCreds(tc testgrpc.TestServiceClient, serviceAccountKeyFile string) {
   315  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   316  	req := &testpb.SimpleRequest{
   317  		ResponseType: testpb.PayloadType_COMPRESSABLE,
   318  		ResponseSize: int32(largeRespSize),
   319  		Payload:      pl,
   320  		FillUsername: true,
   321  	}
   322  	reply, err := tc.UnaryCall(context.Background(), req)
   323  	if err != nil {
   324  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   325  	}
   326  	jsonKey := getServiceAccountJSONKey(serviceAccountKeyFile)
   327  	user := reply.GetUsername()
   328  	if !strings.Contains(string(jsonKey), user) {
   329  		logger.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
   330  	}
   331  }
   332  
   333  // GetToken obtains an OAUTH token from the input.
   334  func GetToken(serviceAccountKeyFile string, oauthScope string) *oauth2.Token {
   335  	jsonKey := getServiceAccountJSONKey(serviceAccountKeyFile)
   336  	config, err := google.JWTConfigFromJSON(jsonKey, oauthScope)
   337  	if err != nil {
   338  		logger.Fatalf("Failed to get the config: %v", err)
   339  	}
   340  	token, err := config.TokenSource(context.Background()).Token()
   341  	if err != nil {
   342  		logger.Fatalf("Failed to get the token: %v", err)
   343  	}
   344  	return token
   345  }
   346  
   347  // DoOauth2TokenCreds performs a unary RPC with OAUTH2 token auth.
   348  func DoOauth2TokenCreds(tc testgrpc.TestServiceClient, serviceAccountKeyFile, oauthScope string) {
   349  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   350  	req := &testpb.SimpleRequest{
   351  		ResponseType:   testpb.PayloadType_COMPRESSABLE,
   352  		ResponseSize:   int32(largeRespSize),
   353  		Payload:        pl,
   354  		FillUsername:   true,
   355  		FillOauthScope: true,
   356  	}
   357  	reply, err := tc.UnaryCall(context.Background(), req)
   358  	if err != nil {
   359  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   360  	}
   361  	jsonKey := getServiceAccountJSONKey(serviceAccountKeyFile)
   362  	user := reply.GetUsername()
   363  	scope := reply.GetOauthScope()
   364  	if !strings.Contains(string(jsonKey), user) {
   365  		logger.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
   366  	}
   367  	if !strings.Contains(oauthScope, scope) {
   368  		logger.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, oauthScope)
   369  	}
   370  }
   371  
   372  // DoPerRPCCreds performs a unary RPC with per RPC OAUTH2 token.
   373  func DoPerRPCCreds(tc testgrpc.TestServiceClient, serviceAccountKeyFile, oauthScope string) {
   374  	jsonKey := getServiceAccountJSONKey(serviceAccountKeyFile)
   375  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   376  	req := &testpb.SimpleRequest{
   377  		ResponseType:   testpb.PayloadType_COMPRESSABLE,
   378  		ResponseSize:   int32(largeRespSize),
   379  		Payload:        pl,
   380  		FillUsername:   true,
   381  		FillOauthScope: true,
   382  	}
   383  	token := GetToken(serviceAccountKeyFile, oauthScope)
   384  	kv := map[string]string{"authorization": token.Type() + " " + token.AccessToken}
   385  	ctx := metadata.NewOutgoingContext(context.Background(), metadata.MD{"authorization": []string{kv["authorization"]}})
   386  	reply, err := tc.UnaryCall(ctx, req)
   387  	if err != nil {
   388  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   389  	}
   390  	user := reply.GetUsername()
   391  	scope := reply.GetOauthScope()
   392  	if !strings.Contains(string(jsonKey), user) {
   393  		logger.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
   394  	}
   395  	if !strings.Contains(oauthScope, scope) {
   396  		logger.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, oauthScope)
   397  	}
   398  }
   399  
   400  // DoGoogleDefaultCredentials performs an unary RPC with google default credentials
   401  func DoGoogleDefaultCredentials(tc testgrpc.TestServiceClient, defaultServiceAccount string) {
   402  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   403  	req := &testpb.SimpleRequest{
   404  		ResponseType:   testpb.PayloadType_COMPRESSABLE,
   405  		ResponseSize:   int32(largeRespSize),
   406  		Payload:        pl,
   407  		FillUsername:   true,
   408  		FillOauthScope: true,
   409  	}
   410  	reply, err := tc.UnaryCall(context.Background(), req)
   411  	if err != nil {
   412  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   413  	}
   414  	if reply.GetUsername() != defaultServiceAccount {
   415  		logger.Fatalf("Got user name %q; wanted %q. ", reply.GetUsername(), defaultServiceAccount)
   416  	}
   417  }
   418  
   419  // DoComputeEngineChannelCredentials performs an unary RPC with compute engine channel credentials
   420  func DoComputeEngineChannelCredentials(tc testgrpc.TestServiceClient, defaultServiceAccount string) {
   421  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   422  	req := &testpb.SimpleRequest{
   423  		ResponseType:   testpb.PayloadType_COMPRESSABLE,
   424  		ResponseSize:   int32(largeRespSize),
   425  		Payload:        pl,
   426  		FillUsername:   true,
   427  		FillOauthScope: true,
   428  	}
   429  	reply, err := tc.UnaryCall(context.Background(), req)
   430  	if err != nil {
   431  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   432  	}
   433  	if reply.GetUsername() != defaultServiceAccount {
   434  		logger.Fatalf("Got user name %q; wanted %q. ", reply.GetUsername(), defaultServiceAccount)
   435  	}
   436  }
   437  
   438  var testMetadata = metadata.MD{
   439  	"key1": []string{"value1"},
   440  	"key2": []string{"value2"},
   441  }
   442  
   443  // DoCancelAfterBegin cancels the RPC after metadata has been sent but before payloads are sent.
   444  func DoCancelAfterBegin(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   445  	ctx, cancel := context.WithCancel(metadata.NewOutgoingContext(context.Background(), testMetadata))
   446  	stream, err := tc.StreamingInputCall(ctx, args...)
   447  	if err != nil {
   448  		logger.Fatalf("%v.StreamingInputCall(_) = _, %v", tc, err)
   449  	}
   450  	cancel()
   451  	_, err = stream.CloseAndRecv()
   452  	if status.Code(err) != codes.Canceled {
   453  		logger.Fatalf("%v.CloseAndRecv() got error code %d, want %d", stream, status.Code(err), codes.Canceled)
   454  	}
   455  }
   456  
   457  // DoCancelAfterFirstResponse cancels the RPC after receiving the first message from the server.
   458  func DoCancelAfterFirstResponse(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   459  	ctx, cancel := context.WithCancel(context.Background())
   460  	stream, err := tc.FullDuplexCall(ctx, args...)
   461  	if err != nil {
   462  		logger.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
   463  	}
   464  	respParam := []*testpb.ResponseParameters{
   465  		{
   466  			Size: 31415,
   467  		},
   468  	}
   469  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, 27182)
   470  	req := &testpb.StreamingOutputCallRequest{
   471  		ResponseType:       testpb.PayloadType_COMPRESSABLE,
   472  		ResponseParameters: respParam,
   473  		Payload:            pl,
   474  	}
   475  	if err := stream.Send(req); err != nil {
   476  		logger.Fatalf("%v has error %v while sending %v", stream, err, req)
   477  	}
   478  	if _, err := stream.Recv(); err != nil {
   479  		logger.Fatalf("%v.Recv() = %v", stream, err)
   480  	}
   481  	cancel()
   482  	if _, err := stream.Recv(); status.Code(err) != codes.Canceled {
   483  		logger.Fatalf("%v compleled with error code %d, want %d", stream, status.Code(err), codes.Canceled)
   484  	}
   485  }
   486  
   487  var (
   488  	initialMetadataValue  = "test_initial_metadata_value"
   489  	trailingMetadataValue = "\x0a\x0b\x0a\x0b\x0a\x0b"
   490  	customMetadata        = metadata.Pairs(
   491  		initialMetadataKey, initialMetadataValue,
   492  		trailingMetadataKey, trailingMetadataValue,
   493  	)
   494  )
   495  
   496  func validateMetadata(header, trailer metadata.MD) {
   497  	if len(header[initialMetadataKey]) != 1 {
   498  		logger.Fatalf("Expected exactly one header from server. Received %d", len(header[initialMetadataKey]))
   499  	}
   500  	if header[initialMetadataKey][0] != initialMetadataValue {
   501  		logger.Fatalf("Got header %s; want %s", header[initialMetadataKey][0], initialMetadataValue)
   502  	}
   503  	if len(trailer[trailingMetadataKey]) != 1 {
   504  		logger.Fatalf("Expected exactly one trailer from server. Received %d", len(trailer[trailingMetadataKey]))
   505  	}
   506  	if trailer[trailingMetadataKey][0] != trailingMetadataValue {
   507  		logger.Fatalf("Got trailer %s; want %s", trailer[trailingMetadataKey][0], trailingMetadataValue)
   508  	}
   509  }
   510  
   511  // DoCustomMetadata checks that metadata is echoed back to the client.
   512  func DoCustomMetadata(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   513  	// Testing with UnaryCall.
   514  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, 1)
   515  	req := &testpb.SimpleRequest{
   516  		ResponseType: testpb.PayloadType_COMPRESSABLE,
   517  		ResponseSize: int32(1),
   518  		Payload:      pl,
   519  	}
   520  	ctx := metadata.NewOutgoingContext(context.Background(), customMetadata)
   521  	var header, trailer metadata.MD
   522  	args = append(args, grpc.Header(&header), grpc.Trailer(&trailer))
   523  	reply, err := tc.UnaryCall(
   524  		ctx,
   525  		req,
   526  		args...,
   527  	)
   528  	if err != nil {
   529  		logger.Fatal("/TestService/UnaryCall RPC failed: ", err)
   530  	}
   531  	t := reply.GetPayload().GetType()
   532  	s := len(reply.GetPayload().GetBody())
   533  	if t != testpb.PayloadType_COMPRESSABLE || s != 1 {
   534  		logger.Fatalf("Got the reply with type %d len %d; want %d, %d", t, s, testpb.PayloadType_COMPRESSABLE, 1)
   535  	}
   536  	validateMetadata(header, trailer)
   537  
   538  	// Testing with FullDuplex.
   539  	stream, err := tc.FullDuplexCall(ctx, args...)
   540  	if err != nil {
   541  		logger.Fatalf("%v.FullDuplexCall(_) = _, %v, want <nil>", tc, err)
   542  	}
   543  	respParam := []*testpb.ResponseParameters{
   544  		{
   545  			Size: 1,
   546  		},
   547  	}
   548  	streamReq := &testpb.StreamingOutputCallRequest{
   549  		ResponseType:       testpb.PayloadType_COMPRESSABLE,
   550  		ResponseParameters: respParam,
   551  		Payload:            pl,
   552  	}
   553  	if err := stream.Send(streamReq); err != nil {
   554  		logger.Fatalf("%v has error %v while sending %v", stream, err, streamReq)
   555  	}
   556  	streamHeader, err := stream.Header()
   557  	if err != nil {
   558  		logger.Fatalf("%v.Header() = %v", stream, err)
   559  	}
   560  	if _, err := stream.Recv(); err != nil {
   561  		logger.Fatalf("%v.Recv() = %v", stream, err)
   562  	}
   563  	if err := stream.CloseSend(); err != nil {
   564  		logger.Fatalf("%v.CloseSend() = %v, want <nil>", stream, err)
   565  	}
   566  	if _, err := stream.Recv(); err != io.EOF {
   567  		logger.Fatalf("%v failed to complete the custom metadata test: %v", stream, err)
   568  	}
   569  	streamTrailer := stream.Trailer()
   570  	validateMetadata(streamHeader, streamTrailer)
   571  }
   572  
   573  // DoStatusCodeAndMessage checks that the status code is propagated back to the client.
   574  func DoStatusCodeAndMessage(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   575  	var code int32 = 2
   576  	msg := "test status message"
   577  	expectedErr := status.Error(codes.Code(code), msg)
   578  	respStatus := &testpb.EchoStatus{
   579  		Code:    code,
   580  		Message: msg,
   581  	}
   582  	// Test UnaryCall.
   583  	req := &testpb.SimpleRequest{
   584  		ResponseStatus: respStatus,
   585  	}
   586  	if _, err := tc.UnaryCall(context.Background(), req, args...); err == nil || err.Error() != expectedErr.Error() {
   587  		logger.Fatalf("%v.UnaryCall(_, %v) = _, %v, want _, %v", tc, req, err, expectedErr)
   588  	}
   589  	// Test FullDuplexCall.
   590  	stream, err := tc.FullDuplexCall(context.Background(), args...)
   591  	if err != nil {
   592  		logger.Fatalf("%v.FullDuplexCall(_) = _, %v, want <nil>", tc, err)
   593  	}
   594  	streamReq := &testpb.StreamingOutputCallRequest{
   595  		ResponseStatus: respStatus,
   596  	}
   597  	if err := stream.Send(streamReq); err != nil {
   598  		logger.Fatalf("%v has error %v while sending %v, want <nil>", stream, err, streamReq)
   599  	}
   600  	if err := stream.CloseSend(); err != nil {
   601  		logger.Fatalf("%v.CloseSend() = %v, want <nil>", stream, err)
   602  	}
   603  	if _, err = stream.Recv(); err.Error() != expectedErr.Error() {
   604  		logger.Fatalf("%v.Recv() returned error %v, want %v", stream, err, expectedErr)
   605  	}
   606  }
   607  
   608  // DoSpecialStatusMessage verifies Unicode and whitespace is correctly processed
   609  // in status message.
   610  func DoSpecialStatusMessage(tc testgrpc.TestServiceClient, args ...grpc.CallOption) {
   611  	const (
   612  		code int32  = 2
   613  		msg  string = "\t\ntest with whitespace\r\nand Unicode BMP ☺ and non-BMP 😈\t\n"
   614  	)
   615  	expectedErr := status.Error(codes.Code(code), msg)
   616  	req := &testpb.SimpleRequest{
   617  		ResponseStatus: &testpb.EchoStatus{
   618  			Code:    code,
   619  			Message: msg,
   620  		},
   621  	}
   622  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   623  	defer cancel()
   624  	if _, err := tc.UnaryCall(ctx, req, args...); err == nil || err.Error() != expectedErr.Error() {
   625  		logger.Fatalf("%v.UnaryCall(_, %v) = _, %v, want _, %v", tc, req, err, expectedErr)
   626  	}
   627  }
   628  
   629  // DoUnimplementedService attempts to call a method from an unimplemented service.
   630  func DoUnimplementedService(tc testgrpc.UnimplementedServiceClient) {
   631  	_, err := tc.UnimplementedCall(context.Background(), &testpb.Empty{})
   632  	if status.Code(err) != codes.Unimplemented {
   633  		logger.Fatalf("%v.UnimplementedCall() = _, %v, want _, %v", tc, status.Code(err), codes.Unimplemented)
   634  	}
   635  }
   636  
   637  // DoUnimplementedMethod attempts to call an unimplemented method.
   638  func DoUnimplementedMethod(cc *grpc.ClientConn) {
   639  	var req, reply proto.Message
   640  	if err := cc.Invoke(context.Background(), "/grpc.testing.TestService/UnimplementedCall", req, reply); err == nil || status.Code(err) != codes.Unimplemented {
   641  		logger.Fatalf("ClientConn.Invoke(_, _, _, _, _) = %v, want error code %s", err, codes.Unimplemented)
   642  	}
   643  }
   644  
   645  // DoPickFirstUnary runs multiple RPCs (rpcCount) and checks that all requests
   646  // are sent to the same backend.
   647  func DoPickFirstUnary(tc testgrpc.TestServiceClient) {
   648  	const rpcCount = 100
   649  
   650  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, 1)
   651  	req := &testpb.SimpleRequest{
   652  		ResponseType: testpb.PayloadType_COMPRESSABLE,
   653  		ResponseSize: int32(1),
   654  		Payload:      pl,
   655  		FillServerId: true,
   656  	}
   657  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   658  	defer cancel()
   659  	var serverID string
   660  	for i := 0; i < rpcCount; i++ {
   661  		resp, err := tc.UnaryCall(ctx, req)
   662  		if err != nil {
   663  			logger.Fatalf("iteration %d, failed to do UnaryCall: %v", i, err)
   664  		}
   665  		id := resp.ServerId
   666  		if id == "" {
   667  			logger.Fatalf("iteration %d, got empty server ID", i)
   668  		}
   669  		if i == 0 {
   670  			serverID = id
   671  			continue
   672  		}
   673  		if serverID != id {
   674  			logger.Fatalf("iteration %d, got different server ids: %q vs %q", i, serverID, id)
   675  		}
   676  	}
   677  }
   678  
   679  func doOneSoakIteration(ctx context.Context, tc testgrpc.TestServiceClient, resetChannel bool, serverAddr string, dopts []grpc.DialOption) (latency time.Duration, err error) {
   680  	start := time.Now()
   681  	client := tc
   682  	if resetChannel {
   683  		var conn *grpc.ClientConn
   684  		conn, err = grpc.Dial(serverAddr, dopts...)
   685  		if err != nil {
   686  			return
   687  		}
   688  		defer conn.Close()
   689  		client = testgrpc.NewTestServiceClient(conn)
   690  	}
   691  	// per test spec, don't include channel shutdown in latency measurement
   692  	defer func() { latency = time.Since(start) }()
   693  	// do a large-unary RPC
   694  	pl := ClientNewPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
   695  	req := &testpb.SimpleRequest{
   696  		ResponseType: testpb.PayloadType_COMPRESSABLE,
   697  		ResponseSize: int32(largeRespSize),
   698  		Payload:      pl,
   699  	}
   700  	var reply *testpb.SimpleResponse
   701  	reply, err = client.UnaryCall(ctx, req)
   702  	if err != nil {
   703  		err = fmt.Errorf("/TestService/UnaryCall RPC failed: %s", err)
   704  		return
   705  	}
   706  	t := reply.GetPayload().GetType()
   707  	s := len(reply.GetPayload().GetBody())
   708  	if t != testpb.PayloadType_COMPRESSABLE || s != largeRespSize {
   709  		err = fmt.Errorf("got the reply with type %d len %d; want %d, %d", t, s, testpb.PayloadType_COMPRESSABLE, largeRespSize)
   710  		return
   711  	}
   712  	return
   713  }
   714  
   715  // DoSoakTest runs large unary RPCs in a loop for a configurable number of times, with configurable failure thresholds.
   716  // If resetChannel is false, then each RPC will be performed on tc. Otherwise, each RPC will be performed on a new
   717  // stub that is created with the provided server address and dial options.
   718  func DoSoakTest(tc testgrpc.TestServiceClient, serverAddr string, dopts []grpc.DialOption, resetChannel bool, soakIterations int, maxFailures int, perIterationMaxAcceptableLatency time.Duration, overallDeadline time.Time) {
   719  	start := time.Now()
   720  	ctx, cancel := context.WithDeadline(context.Background(), overallDeadline)
   721  	defer cancel()
   722  	iterationsDone := 0
   723  	totalFailures := 0
   724  	hopts := stats.HistogramOptions{
   725  		NumBuckets:     20,
   726  		GrowthFactor:   1,
   727  		BaseBucketSize: 1,
   728  		MinValue:       0,
   729  	}
   730  	h := stats.NewHistogram(hopts)
   731  	for i := 0; i < soakIterations; i++ {
   732  		if time.Now().After(overallDeadline) {
   733  			break
   734  		}
   735  		iterationsDone++
   736  		latency, err := doOneSoakIteration(ctx, tc, resetChannel, serverAddr, dopts)
   737  		latencyMs := int64(latency / time.Millisecond)
   738  		h.Add(latencyMs)
   739  		if err != nil {
   740  			totalFailures++
   741  			fmt.Fprintf(os.Stderr, "soak iteration: %d elapsed_ms: %d failed: %s\n", i, latencyMs, err)
   742  			continue
   743  		}
   744  		if latency > perIterationMaxAcceptableLatency {
   745  			totalFailures++
   746  			fmt.Fprintf(os.Stderr, "soak iteration: %d elapsed_ms: %d exceeds max acceptable latency: %d\n", i, latencyMs, perIterationMaxAcceptableLatency.Milliseconds())
   747  			continue
   748  		}
   749  		fmt.Fprintf(os.Stderr, "soak iteration: %d elapsed_ms: %d succeeded\n", i, latencyMs)
   750  	}
   751  	var b bytes.Buffer
   752  	h.Print(&b)
   753  	fmt.Fprintln(os.Stderr, "Histogram of per-iteration latencies in milliseconds:")
   754  	fmt.Fprintln(os.Stderr, b.String())
   755  	fmt.Fprintf(os.Stderr, "soak test ran: %d / %d iterations. total failures: %d. max failures threshold: %d. See breakdown above for which iterations succeeded, failed, and why for more info.\n", iterationsDone, soakIterations, totalFailures, maxFailures)
   756  	if iterationsDone < soakIterations {
   757  		logger.Fatalf("soak test consumed all %f seconds of time and quit early, only having ran %d out of desired %d iterations.", overallDeadline.Sub(start).Seconds(), iterationsDone, soakIterations)
   758  	}
   759  	if totalFailures > maxFailures {
   760  		logger.Fatalf("soak test total failures: %d exceeds max failures threshold: %d.", totalFailures, maxFailures)
   761  	}
   762  }
   763  
   764  type testServer struct {
   765  	testgrpc.UnimplementedTestServiceServer
   766  }
   767  
   768  // NewTestServer creates a test server for test service.
   769  func NewTestServer() testgrpc.TestServiceServer {
   770  	return &testServer{}
   771  }
   772  
   773  func (s *testServer) EmptyCall(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
   774  	return new(testpb.Empty), nil
   775  }
   776  
   777  func serverNewPayload(t testpb.PayloadType, size int32) (*testpb.Payload, error) {
   778  	if size < 0 {
   779  		return nil, fmt.Errorf("requested a response with invalid length %d", size)
   780  	}
   781  	body := make([]byte, size)
   782  	switch t {
   783  	case testpb.PayloadType_COMPRESSABLE:
   784  	default:
   785  		return nil, fmt.Errorf("unsupported payload type: %d", t)
   786  	}
   787  	return &testpb.Payload{
   788  		Type: t,
   789  		Body: body,
   790  	}, nil
   791  }
   792  
   793  func (s *testServer) UnaryCall(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
   794  	st := in.GetResponseStatus()
   795  	if md, ok := metadata.FromIncomingContext(ctx); ok {
   796  		if initialMetadata, ok := md[initialMetadataKey]; ok {
   797  			header := metadata.Pairs(initialMetadataKey, initialMetadata[0])
   798  			grpc.SendHeader(ctx, header)
   799  		}
   800  		if trailingMetadata, ok := md[trailingMetadataKey]; ok {
   801  			trailer := metadata.Pairs(trailingMetadataKey, trailingMetadata[0])
   802  			grpc.SetTrailer(ctx, trailer)
   803  		}
   804  	}
   805  	if st != nil && st.Code != 0 {
   806  		return nil, status.Error(codes.Code(st.Code), st.Message)
   807  	}
   808  	pl, err := serverNewPayload(in.GetResponseType(), in.GetResponseSize())
   809  	if err != nil {
   810  		return nil, err
   811  	}
   812  	return &testpb.SimpleResponse{
   813  		Payload: pl,
   814  	}, nil
   815  }
   816  
   817  func (s *testServer) StreamingOutputCall(args *testpb.StreamingOutputCallRequest, stream testgrpc.TestService_StreamingOutputCallServer) error {
   818  	cs := args.GetResponseParameters()
   819  	for _, c := range cs {
   820  		if us := c.GetIntervalUs(); us > 0 {
   821  			time.Sleep(time.Duration(us) * time.Microsecond)
   822  		}
   823  		pl, err := serverNewPayload(args.GetResponseType(), c.GetSize())
   824  		if err != nil {
   825  			return err
   826  		}
   827  		if err := stream.Send(&testpb.StreamingOutputCallResponse{
   828  			Payload: pl,
   829  		}); err != nil {
   830  			return err
   831  		}
   832  	}
   833  	return nil
   834  }
   835  
   836  func (s *testServer) StreamingInputCall(stream testgrpc.TestService_StreamingInputCallServer) error {
   837  	var sum int
   838  	for {
   839  		in, err := stream.Recv()
   840  		if err == io.EOF {
   841  			return stream.SendAndClose(&testpb.StreamingInputCallResponse{
   842  				AggregatedPayloadSize: int32(sum),
   843  			})
   844  		}
   845  		if err != nil {
   846  			return err
   847  		}
   848  		p := in.GetPayload().GetBody()
   849  		sum += len(p)
   850  	}
   851  }
   852  
   853  func (s *testServer) FullDuplexCall(stream testgrpc.TestService_FullDuplexCallServer) error {
   854  	if md, ok := metadata.FromIncomingContext(stream.Context()); ok {
   855  		if initialMetadata, ok := md[initialMetadataKey]; ok {
   856  			header := metadata.Pairs(initialMetadataKey, initialMetadata[0])
   857  			stream.SendHeader(header)
   858  		}
   859  		if trailingMetadata, ok := md[trailingMetadataKey]; ok {
   860  			trailer := metadata.Pairs(trailingMetadataKey, trailingMetadata[0])
   861  			stream.SetTrailer(trailer)
   862  		}
   863  	}
   864  	for {
   865  		in, err := stream.Recv()
   866  		if err == io.EOF {
   867  			// read done.
   868  			return nil
   869  		}
   870  		if err != nil {
   871  			return err
   872  		}
   873  		st := in.GetResponseStatus()
   874  		if st != nil && st.Code != 0 {
   875  			return status.Error(codes.Code(st.Code), st.Message)
   876  		}
   877  		cs := in.GetResponseParameters()
   878  		for _, c := range cs {
   879  			if us := c.GetIntervalUs(); us > 0 {
   880  				time.Sleep(time.Duration(us) * time.Microsecond)
   881  			}
   882  			pl, err := serverNewPayload(in.GetResponseType(), c.GetSize())
   883  			if err != nil {
   884  				return err
   885  			}
   886  			if err := stream.Send(&testpb.StreamingOutputCallResponse{
   887  				Payload: pl,
   888  			}); err != nil {
   889  				return err
   890  			}
   891  		}
   892  	}
   893  }
   894  
   895  func (s *testServer) HalfDuplexCall(stream testgrpc.TestService_HalfDuplexCallServer) error {
   896  	var msgBuf []*testpb.StreamingOutputCallRequest
   897  	for {
   898  		in, err := stream.Recv()
   899  		if err == io.EOF {
   900  			// read done.
   901  			break
   902  		}
   903  		if err != nil {
   904  			return err
   905  		}
   906  		msgBuf = append(msgBuf, in)
   907  	}
   908  	for _, m := range msgBuf {
   909  		cs := m.GetResponseParameters()
   910  		for _, c := range cs {
   911  			if us := c.GetIntervalUs(); us > 0 {
   912  				time.Sleep(time.Duration(us) * time.Microsecond)
   913  			}
   914  			pl, err := serverNewPayload(m.GetResponseType(), c.GetSize())
   915  			if err != nil {
   916  				return err
   917  			}
   918  			if err := stream.Send(&testpb.StreamingOutputCallResponse{
   919  				Payload: pl,
   920  			}); err != nil {
   921  				return err
   922  			}
   923  		}
   924  	}
   925  	return nil
   926  }