Merge pull request #39 from k1LoW/enable-health

Add EnableHealthCheck option

grpcstub.go

@@ -26,6 +26,8 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/reflection"
 	"google.golang.org/grpc/status"
@@ -71,16 +73,17 @@ func NewResponse() *Response {
 }
 
 type Server struct {
-	matchers []*matcher
-	fds      []*desc.FileDescriptor
-	listener net.Listener
-	server   *grpc.Server
-	tlsc     *tls.Config
-	cacert   []byte
-	cc       *grpc.ClientConn
-	requests []*Request
-	t        *testing.T
-	mu       sync.RWMutex
+	matchers    []*matcher
+	fds         []*desc.FileDescriptor
+	listener    net.Listener
+	server      *grpc.Server
+	tlsc        *tls.Config
+	cacert      []byte
+	cc          *grpc.ClientConn
+	requests    []*Request
+	healthCheck bool
+	t           *testing.T
+	mu          sync.RWMutex
 }
 
 type matcher struct {
@@ -109,8 +112,9 @@ func NewServer(t *testing.T, protopath string, opts ...Option) *Server {
 		return nil
 	}
 	s := &Server{
-		fds: fds,
-		t:   t,
+		fds:         fds,
+		t:           t,
+		healthCheck: c.healthCheck,
 	}
 	if c.useTLS {
 		certificate, err := tls.X509KeyPair(c.cert, c.key)
@@ -396,6 +400,7 @@ func (m *matcher) Status(s *status.Status) *matcher {
 func (s *Server) Requests() []*Request {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
+	fmt.Printf("%v\n", s.requests)
 	return s.requests
 }
 
@@ -416,6 +421,11 @@ func (s *Server) registerServer() {
 	for _, sd := range sds {
 		s.server.RegisterService(sd, nil)
 	}
+	if s.healthCheck {
+		healthSrv := health.NewServer()
+		healthpb.RegisterHealthServer(s.server, healthSrv)
+		healthSrv.SetServingStatus("grpcstub", healthpb.HealthCheckResponse_SERVING)
+	}
 }
 
 func (s *Server) createServiceDesc(sd *desc.ServiceDescriptor) *grpc.ServiceDesc {
@@ -477,6 +487,7 @@ func (s *Server) createUnaryHandler(md *desc.MethodDescriptor) func(srv interfac
 		}
 		s.mu.Lock()
 		s.requests = append(s.requests, r)
+		fmt.Printf("s.requests: %v\n", s.requests)
 		s.mu.Unlock()
 
 		var mes proto.Message
@@ -567,6 +578,7 @@ func (s *Server) createServerStreamingHandler(md *desc.MethodDescriptor) func(sr
 		}
 		s.mu.Lock()
 		s.requests = append(s.requests, r)
+		fmt.Printf("s.requests: %v\n", s.requests)
 		s.mu.Unlock()
 		for _, m := range s.matchers {
 			match := true

grpcstub_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/k1LoW/grpcstub/testdata/routeguide"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	healthpb "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/metadata"
 	rpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
 	"google.golang.org/grpc/status"
@@ -789,3 +790,40 @@ func TestTLSServer(t *testing.T) {
 		}
 	}
 }
+
+func TestHealthCheck(t *testing.T) {
+	tests := []struct {
+		enable  bool
+		wantErr bool
+	}{
+		{true, false},
+		{false, true},
+	}
+	ctx := context.Background()
+	for i, tt := range tests {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			var ts *Server
+			if tt.enable {
+				ts = NewServer(t, "testdata/*.proto", EnableHealthCheck())
+			} else {
+				ts = NewServer(t, "testdata/*.proto")
+			}
+			t.Cleanup(func() {
+				ts.Close()
+			})
+			client := healthpb.NewHealthClient(ts.ClientConn())
+			_, err := client.Check(ctx, &healthpb.HealthCheckRequest{
+				Service: "grpcstub",
+			})
+			if err != nil {
+				if !tt.wantErr {
+					t.Errorf("got error: %s", err)
+				}
+				return
+			}
+			if tt.wantErr {
+				t.Error("want error")
+			}
+		})
+	}
+}

option.go

@@ -13,6 +13,7 @@ type config struct {
 	importPaths       []string
 	useTLS            bool
 	cacert, cert, key []byte
+	healthCheck       bool
 }
 
 type Option func(*config) error
@@ -93,6 +94,14 @@ func UseTLS(cacert, cert, key []byte) Option {
 	}
 }
 
+// EnableHealthCheck enable grpc.health.v1
+func EnableHealthCheck() Option {
+	return func(c *config) error {
+		c.healthCheck = true
+		return nil
+	}
+}
+
 func unique(in []string) []string {
 	u := []string{}
 	m := map[string]struct{}{}