mockserver.go 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. // Copyright 2018 The etcd 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. package mockserver
  15. import (
  16. "context"
  17. "fmt"
  18. "net"
  19. pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
  20. "google.golang.org/grpc"
  21. )
  22. // MockServer provides a mocked out grpc server of the etcdserver interface.
  23. type MockServer struct {
  24. GrpcServer *grpc.Server
  25. Address string
  26. }
  27. // MockServers provides a cluster of mocket out gprc servers of the etcdserver interface.
  28. type MockServers []*MockServer
  29. // StartMockServers creates the desired count of mock servers
  30. // and starts them.
  31. func StartMockServers(count int) (svrs MockServers, err error) {
  32. svrs = make(MockServers, count)
  33. defer func() {
  34. if err != nil {
  35. svrs.Stop()
  36. }
  37. }()
  38. for i := 0; i < count; i++ {
  39. listener, err := net.Listen("tcp", "localhost:0")
  40. if err != nil {
  41. return nil, fmt.Errorf("failed to listen %v", err)
  42. }
  43. svr := grpc.NewServer()
  44. pb.RegisterKVServer(svr, &mockKVServer{})
  45. svrs[i] = &MockServer{GrpcServer: svr, Address: listener.Addr().String()}
  46. go func(svr *grpc.Server, l net.Listener) {
  47. svr.Serve(l)
  48. }(svr, listener)
  49. }
  50. return svrs, nil
  51. }
  52. // Stop stops the mock server, immediately closing all open connections and listeners.
  53. func (svrs MockServers) Stop() {
  54. for _, svr := range svrs {
  55. svr.GrpcServer.Stop()
  56. }
  57. }
  58. type mockKVServer struct{}
  59. func (m *mockKVServer) Range(context.Context, *pb.RangeRequest) (*pb.RangeResponse, error) {
  60. return &pb.RangeResponse{}, nil
  61. }
  62. func (m *mockKVServer) Put(context.Context, *pb.PutRequest) (*pb.PutResponse, error) {
  63. return &pb.PutResponse{}, nil
  64. }
  65. func (m *mockKVServer) DeleteRange(context.Context, *pb.DeleteRangeRequest) (*pb.DeleteRangeResponse, error) {
  66. return &pb.DeleteRangeResponse{}, nil
  67. }
  68. func (m *mockKVServer) Txn(context.Context, *pb.TxnRequest) (*pb.TxnResponse, error) {
  69. return &pb.TxnResponse{}, nil
  70. }
  71. func (m *mockKVServer) Compact(context.Context, *pb.CompactionRequest) (*pb.CompactionResponse, error) {
  72. return &pb.CompactionResponse{}, nil
  73. }