stub_server.go

Documentation: go.etcd.io/etcd/pkg/v3/grpc_testing

     1  package grpc_testing
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net"
     7  
     8  	"google.golang.org/grpc"
     9  	testpb "google.golang.org/grpc/interop/grpc_testing"
    10  )
    11  
    12  // StubServer is borrowed from the interal package of grpc-go.
    13  // See https://github.com/grpc/grpc-go/blob/master/internal/stubserver/stubserver.go
    14  // Since it cannot be imported directly, we have to copy and paste it here,
    15  // and useless code for our testing is removed.
    16  
    17  // StubServer is a server that is easy to customize within individual test
    18  // cases.
    19  type StubServer struct {
    20  	testService testpb.TestServiceServer
    21  
    22  	// Network and Address are parameters for Listen. Defaults will be used if these are empty before Start.
    23  	Network string
    24  	Address string
    25  
    26  	s *grpc.Server
    27  
    28  	cleanups []func() // Lambdas executed in Stop(); populated by Start().
    29  }
    30  
    31  func New(testService testpb.TestServiceServer) *StubServer {
    32  	return &StubServer{testService: testService}
    33  }
    34  
    35  // Start starts the server and creates a client connected to it.
    36  func (ss *StubServer) Start(sopts []grpc.ServerOption, dopts ...grpc.DialOption) error {
    37  	if ss.Network == "" {
    38  		ss.Network = "tcp"
    39  	}
    40  	if ss.Address == "" {
    41  		ss.Address = "localhost:0"
    42  	}
    43  
    44  	lis, err := net.Listen(ss.Network, ss.Address)
    45  	if err != nil {
    46  		return fmt.Errorf("net.Listen(%q, %q) = %v", ss.Network, ss.Address, err)
    47  	}
    48  	ss.Address = lis.Addr().String()
    49  	ss.cleanups = append(ss.cleanups, func() { lis.Close() })
    50  
    51  	s := grpc.NewServer(sopts...)
    52  	testpb.RegisterTestServiceServer(s, ss.testService)
    53  	go s.Serve(lis)
    54  	ss.cleanups = append(ss.cleanups, s.Stop)
    55  	ss.s = s
    56  
    57  	return nil
    58  }
    59  
    60  // Stop stops ss and cleans up all resources it consumed.
    61  func (ss *StubServer) Stop() {
    62  	for i := len(ss.cleanups) - 1; i >= 0; i-- {
    63  		ss.cleanups[i]()
    64  	}
    65  }
    66  
    67  // Addr gets the address the server listening on.
    68  func (ss *StubServer) Addr() string {
    69  	return ss.Address
    70  }
    71  
    72  type dummyStubServer struct {
    73  	testpb.UnimplementedTestServiceServer
    74  	body []byte
    75  }
    76  
    77  func (d dummyStubServer) UnaryCall(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
    78  	return &testpb.SimpleResponse{
    79  		Payload: &testpb.Payload{
    80  			Type: testpb.PayloadType_COMPRESSABLE,
    81  			Body: d.body,
    82  		},
    83  	}, nil
    84  }
    85  
    86  // NewDummyStubServer creates a simple test server that serves Unary calls with
    87  // responses with the given payload.
    88  func NewDummyStubServer(body []byte) *StubServer {
    89  	return New(dummyStubServer{body: body})
    90  }
    91
View as plain text