modify the status map type

guoyu07 · Jul 24, 2015 · e6a0e51 · e6a0e51
1 parent a09d93d
commit e6a0e51
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/health/health.go b/health/health.go
@@ -14,12 +14,12 @@ import (
 type HealthServer struct {
 	mu sync.Mutex
 	// statusMap stores the serving status of the services this HealthServer monitors.
-	statusMap map[string]int32
+	statusMap map[string]healthpb.HealthCheckResponse_ServingStatus
 }
 
 func NewHealthServer() *HealthServer {
 	return &HealthServer{
-		statusMap: make(map[string]int32),
+		statusMap: make(map[string]healthpb.HealthCheckResponse_ServingStatus),
 	}
 }
 
@@ -29,15 +29,15 @@ func (s *HealthServer) Check(ctx context.Context, in *healthpb.HealthCheckReques
 	defer s.mu.Unlock()
 	if status, ok := s.statusMap[service]; ok {
 		return &healthpb.HealthCheckResponse{
-			Status: healthpb.HealthCheckResponse_ServingStatus(status),
+			Status: status,
 		}, nil
 	}
 	return nil, grpc.Errorf(codes.NotFound, "unknown service")
 }
 
 // SetServingStatus is called when need to reset the serving status of a service
 // or insert a new service entry into the statusMap.
-func (s *HealthServer) SetServingStatus(host string, service string, status int32) {
+func (s *HealthServer) SetServingStatus(host string, service string, status healthpb.HealthCheckResponse_ServingStatus) {
 	service = host + ":" + service
 	s.mu.Lock()
 	s.statusMap[service] = status

diff --git a/test/end2end_test.go b/test/end2end_test.go
@@ -430,15 +430,15 @@ func testHealthCheckServingStatus(t *testing.T, e env) {
 	if _, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck"); err != grpc.Errorf(codes.NotFound, "unknown service") {
 		t.Fatalf("HealthCheck/Check(_, _) = _, %v, want _, error code %d", err, codes.NotFound)
 	}
-	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", int32(healthpb.HealthCheckResponse_SERVING))
+	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", healthpb.HealthCheckResponse_SERVING)
 	out, err := healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck")
 	if err != nil {
 		t.Fatalf("HealthCheck/Check(_, _) = _, %v, want _, <nil>", err)
 	}
 	if out.Status != healthpb.HealthCheckResponse_SERVING {
 		t.Fatalf("Got the serving status %v, want SERVING", out.Status)
 	}
-	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", int32(healthpb.HealthCheckResponse_NOT_SERVING))
+	hs.SetServingStatus("", "grpc.health.v1alpha.HealthCheck", healthpb.HealthCheckResponse_NOT_SERVING)
 	out, err = healthCheck(1*time.Second, cc, "grpc.health.v1alpha.HealthCheck")
 	if err != nil {
 		t.Fatalf("HealthCheck/Check(_, _) = _, %v, want _, <nil>", err)