gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/go-grpc-middleware/testing/interceptor_suite.go (about)

     1  // Copyright 2016 Michal Witkowski. All Rights Reserved.
     2  // See LICENSE for licensing terms.
     3  
     4  package grpc_testing
     5  
     6  import (
     7  	"net"
     8  	"time"
     9  
    10  	"flag"
    11  	"path"
    12  	"runtime"
    13  
    14  	pb_testproto "gitee.com/ks-custle/core-gm/go-grpc-middleware/testing/testproto"
    15  	"gitee.com/ks-custle/core-gm/grpc"
    16  	"gitee.com/ks-custle/core-gm/grpc/credentials"
    17  	"gitee.com/ks-custle/core-gm/net/context"
    18  	"github.com/stretchr/testify/require"
    19  	"github.com/stretchr/testify/suite"
    20  )
    21  
    22  var (
    23  	flagTls = flag.Bool("use_tls", true, "whether all gRPC middleware tests should use tls")
    24  )
    25  
    26  func getTestingCertsPath() string {
    27  	_, callerPath, _, _ := runtime.Caller(0)
    28  	return path.Join(path.Dir(callerPath), "certs")
    29  }
    30  
    31  // InterceptorTestSuite is a testify/Suite that starts a gRPC PingService server and a client.
    32  type InterceptorTestSuite struct {
    33  	suite.Suite
    34  
    35  	TestService pb_testproto.TestServiceServer
    36  	ServerOpts  []grpc.ServerOption
    37  	ClientOpts  []grpc.DialOption
    38  
    39  	serverAddr     string
    40  	ServerListener net.Listener
    41  	Server         *grpc.Server
    42  	clientConn     *grpc.ClientConn
    43  	Client         pb_testproto.TestServiceClient
    44  
    45  	restartServerWithDelayedStart chan time.Duration
    46  	serverRunning                 chan bool
    47  }
    48  
    49  func (s *InterceptorTestSuite) SetupSuite() {
    50  	s.restartServerWithDelayedStart = make(chan time.Duration)
    51  	s.serverRunning = make(chan bool)
    52  
    53  	s.serverAddr = "127.0.0.1:0"
    54  
    55  	go func() {
    56  		for {
    57  			var err error
    58  			s.ServerListener, err = net.Listen("tcp", s.serverAddr)
    59  			s.serverAddr = s.ServerListener.Addr().String()
    60  			require.NoError(s.T(), err, "must be able to allocate a port for serverListener")
    61  			if *flagTls {
    62  				creds, err := credentials.NewServerTLSFromFile(
    63  					path.Join(getTestingCertsPath(), "localhost.crt"),
    64  					path.Join(getTestingCertsPath(), "localhost.key"),
    65  				)
    66  				require.NoError(s.T(), err, "failed reading server credentials for localhost.crt")
    67  				s.ServerOpts = append(s.ServerOpts, grpc.Creds(creds))
    68  			}
    69  			// This is the point where we hook up the interceptor
    70  			s.Server = grpc.NewServer(s.ServerOpts...)
    71  			// Crete a service of the instantiator hasn't provided one.
    72  			if s.TestService == nil {
    73  				s.TestService = &TestPingService{T: s.T()}
    74  			}
    75  			pb_testproto.RegisterTestServiceServer(s.Server, s.TestService)
    76  
    77  			go func() {
    78  				s.Server.Serve(s.ServerListener)
    79  			}()
    80  			if s.Client == nil {
    81  				s.Client = s.NewClient(s.ClientOpts...)
    82  			}
    83  
    84  			s.serverRunning <- true
    85  
    86  			d := <-s.restartServerWithDelayedStart
    87  			s.Server.Stop()
    88  			time.Sleep(d)
    89  		}
    90  	}()
    91  
    92  	<-s.serverRunning
    93  }
    94  
    95  func (s *InterceptorTestSuite) RestartServer(delayedStart time.Duration) <-chan bool {
    96  	s.restartServerWithDelayedStart <- delayedStart
    97  	time.Sleep(10 * time.Millisecond)
    98  	return s.serverRunning
    99  }
   100  
   101  func (s *InterceptorTestSuite) NewClient(dialOpts ...grpc.DialOption) pb_testproto.TestServiceClient {
   102  	newDialOpts := append(dialOpts, grpc.WithBlock(), grpc.WithTimeout(2*time.Second))
   103  	if *flagTls {
   104  		creds, err := credentials.NewClientTLSFromFile(
   105  			path.Join(getTestingCertsPath(), "localhost.crt"), "localhost")
   106  		require.NoError(s.T(), err, "failed reading client credentials for localhost.crt")
   107  		newDialOpts = append(newDialOpts, grpc.WithTransportCredentials(creds))
   108  	} else {
   109  		newDialOpts = append(newDialOpts, grpc.WithInsecure())
   110  	}
   111  	clientConn, err := grpc.Dial(s.ServerAddr(), newDialOpts...)
   112  	require.NoError(s.T(), err, "must not error on client Dial")
   113  	return pb_testproto.NewTestServiceClient(clientConn)
   114  }
   115  
   116  func (s *InterceptorTestSuite) ServerAddr() string {
   117  	return s.serverAddr
   118  }
   119  
   120  func (s *InterceptorTestSuite) SimpleCtx() context.Context {
   121  	ctx, _ := context.WithTimeout(context.TODO(), 2*time.Second)
   122  	return ctx
   123  }
   124  
   125  func (s *InterceptorTestSuite) DeadlineCtx(deadline time.Time) context.Context {
   126  	ctx, _ := context.WithDeadline(context.TODO(), deadline)
   127  	return ctx
   128  }
   129  
   130  func (s *InterceptorTestSuite) TearDownSuite() {
   131  	time.Sleep(10 * time.Millisecond)
   132  	if s.ServerListener != nil {
   133  		s.Server.GracefulStop()
   134  		s.T().Logf("stopped grpc.Server at: %v", s.ServerAddr())
   135  		s.ServerListener.Close()
   136  
   137  	}
   138  	if s.clientConn != nil {
   139  		s.clientConn.Close()
   140  	}
   141  }