From 701e5df79efadd4245d955e6b2dc89e83eb57171 Mon Sep 17 00:00:00 2001
From: Ying Li <ying.li@docker.com>
Date: Wed, 7 Oct 2015 18:12:51 -0700
Subject: [PATCH] Change the rpc Server objects to take a function that can
 check health

Signed-off-by: Ying Li <ying.li@docker.com>
---
 signer/api/rpc_api.go      | 20 ++++++++++++++++++++
 signer/api/rpc_api_test.go | 13 +++++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/signer/api/rpc_api.go b/signer/api/rpc_api.go
index de1090c30e..626dabe29a 100644
--- a/signer/api/rpc_api.go
+++ b/signer/api/rpc_api.go
@@ -18,11 +18,13 @@ import (
 //KeyManagementServer implements the KeyManagementServer grpc interface
 type KeyManagementServer struct {
 	CryptoServices signer.CryptoServiceIndex
+    HealthChecker func() map[string]string
 }
 
 //SignerServer implements the SignerServer grpc interface
 type SignerServer struct {
 	CryptoServices signer.CryptoServiceIndex
+    HealthChecker func() map[string]string
 }
 
 //CreateKey returns a PublicKey created using KeyManagementServer's SigningService
@@ -106,6 +108,15 @@ func (s *KeyManagementServer) GetKeyInfo(ctx context.Context, keyID *pb.KeyID) (
 	}, nil
 }
 
+//CheckHealth returns the HealthStatus with the service
+func (s *KeyManagementServer) CheckHealth(ctx context.Context) (*pb.HealthStatus, error) {
+
+    logger.Debug("CheckHealth: Returning HealthStatus for KeyManagementServer")
+    return &pb.HealthStatus{
+        Status: s.HealthChecker(),
+    }, nil
+}
+
 //Sign signs a message and returns the signature using a private key associate with the KeyID from the SignatureRequest
 func (s *SignerServer) Sign(ctx context.Context, sr *pb.SignatureRequest) (*pb.Signature, error) {
 	tufKey, service, err := FindKeyByID(s.CryptoServices, sr.KeyID)
@@ -136,3 +147,12 @@ func (s *SignerServer) Sign(ctx context.Context, sr *pb.SignatureRequest) (*pb.S
 
 	return signature, nil
 }
+
+//CheckHealth returns the HealthStatus with the service
+func (s *SignerServer) CheckHealth(ctx context.Context) (*pb.HealthStatus, error) {
+
+    logger.Debug("CheckHealth: Returning HealthStatus for SignerServer")
+    return &pb.HealthStatus{
+        Status: s.HealthChecker(),
+    }, nil
+}
diff --git a/signer/api/rpc_api_test.go b/signer/api/rpc_api_test.go
index 4021c7dae9..3f2c2fed4e 100644
--- a/signer/api/rpc_api_test.go
+++ b/signer/api/rpc_api_test.go
@@ -26,6 +26,10 @@ var (
 	grpcServer *grpc.Server
 	void       *pb.Void
 	pr         passphrase.Retriever
+    health=    map[string]string {
+        "db": "ok",
+        "other": "not ok",
+    }
 )
 
 func init() {
@@ -52,8 +56,13 @@ func init() {
 	if err != nil {
 		log.Fatalf("fail to dial: %v", err)
 	}
-	kmClient = pb.NewKeyManagementClient(conn)
-	sClient = pb.NewSignerClient(conn)
+
+    fakeHealth := func() map[string]string {
+        return health
+    }
+
+	kmClient = pb.NewKeyManagementClient(conn, fakeHealth)
+	sClient = pb.NewSignerClient(conn, fakeHealth)
 }
 
 func TestDeleteKeyHandlerReturnsNotFoundWithNonexistentKey(t *testing.T) {