github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/milevadb-server/einsteindb/mock_einsteindb_service_test.go (about)

     1  package einsteindb
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"time"
     7  
     8  	"github.com/whtcorpsinc/ekvproto/pkg/einsteindbpb"
     9  	"github.com/whtcorpsinc/milevadb/soliton/logutil"
    10  	"go.uber.org/zap"
    11  	"google.golang.org/grpc"
    12  )
    13  
    14  type server struct {
    15  	einsteindbpb.EinsteinDBServer
    16  }
    17  
    18  func (s *server) BatchCommands(ss einsteindbpb.EinsteinDB_BatchCommandsServer) error {
    19  	for {
    20  		req, err := ss.Recv()
    21  		if err != nil {
    22  			logutil.BgLogger().Error("batch commands receive fail", zap.Error(err))
    23  			return err
    24  		}
    25  
    26  		responses := make([]*einsteindbpb.BatchCommandsResponse_Response, 0, len(req.GetRequestIds()))
    27  		for i := 0; i < len(req.GetRequestIds()); i++ {
    28  			responses = append(responses, &einsteindbpb.BatchCommandsResponse_Response{
    29  				Cmd: &einsteindbpb.BatchCommandsResponse_Response_Empty{
    30  					Empty: &einsteindbpb.BatchCommandsEmptyResponse{},
    31  				},
    32  			})
    33  		}
    34  
    35  		err = ss.Send(&einsteindbpb.BatchCommandsResponse{
    36  			Responses:  responses,
    37  			RequestIds: req.GetRequestIds(),
    38  		})
    39  		if err != nil {
    40  			logutil.BgLogger().Error("batch commands send fail", zap.Error(err))
    41  			return err
    42  		}
    43  	}
    44  }
    45  
    46  // Try to start a gRPC server and retrun the server instance and binded port.
    47  func startMockEinsteinDBService() (*grpc.Server, int) {
    48  	port := -1
    49  	lis, err := net.Listen("tcp", fmt.Sprintf("%s:%d", "127.0.0.1", 0))
    50  	if err != nil {
    51  		logutil.BgLogger().Error("can't listen", zap.Error(err))
    52  		logutil.BgLogger().Error("can't start mock einsteindb service because no available ports")
    53  		return nil, port
    54  	}
    55  	port = lis.Addr().(*net.TCPAddr).Port
    56  	s := grpc.NewServer(grpc.ConnectionTimeout(time.Minute))
    57  	einsteindbpb.RegisterEinsteinDBServer(s, &server{})
    58  	go func() {
    59  		if err = s.Serve(lis); err != nil {
    60  			logutil.BgLogger().Error(
    61  				"can't serve gRPC requests",
    62  				zap.Error(err),
    63  			)
    64  		}
    65  	}()
    66  	return s, port
    67  }