From 97c5d53b8deb904c8cccfecd58fa8fe04e5ad654 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Fri, 9 Jan 2015 18:30:29 -0800
Subject: [PATCH] consul: checks get stored properly

---
 consul/catalog_endpoint.go |  1 +
 consul/state_store.go      |  7 ++++++-
 consul/state_store_test.go | 11 ++++++++++-
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/consul/catalog_endpoint.go b/consul/catalog_endpoint.go
index 0cb2c86275f3..b39c19f2bd7a 100644
--- a/consul/catalog_endpoint.go
+++ b/consul/catalog_endpoint.go
@@ -54,6 +54,7 @@ func (c *Catalog) Register(args *structs.RegisterRequest, reply *struct{}) error
 
 	if args.Check != nil {
 		args.Checks = append(args.Checks, args.Check)
+		args.Check = nil
 	}
 	for _, check := range args.Checks {
 		if check.CheckID == "" && check.Name != "" {
diff --git a/consul/state_store.go b/consul/state_store.go
index 9cf88b16461e..5813a38081e9 100644
--- a/consul/state_store.go
+++ b/consul/state_store.go
@@ -497,12 +497,17 @@ func (s *StateStore) EnsureRegistration(index uint64, req *structs.RegisterReque
 		}
 	}
 
-	// Ensure the check if provided
+	// Ensure the check(s), if provided
 	if req.Check != nil {
 		if err := s.ensureCheckTxn(index, req.Check, tx); err != nil {
 			return err
 		}
 	}
+	for _, check := range req.Checks {
+		if err := s.ensureCheckTxn(index, check, tx); err != nil {
+			return err
+		}
+	}
 
 	// Commit as one unit
 	return tx.Commit()
diff --git a/consul/state_store_test.go b/consul/state_store_test.go
index b456c19e97e7..306b73d3c36b 100644
--- a/consul/state_store_test.go
+++ b/consul/state_store_test.go
@@ -32,6 +32,15 @@ func TestEnsureRegistration(t *testing.T) {
 			Status:    structs.HealthPassing,
 			ServiceID: "api",
 		},
+		Checks: structs.HealthChecks{
+			&structs.HealthCheck{
+				Node:      "foo",
+				CheckID:   "api-cache",
+				Name:      "Can cache stuff",
+				Status:    structs.HealthPassing,
+				ServiceID: "api",
+			},
+		},
 	}
 
 	if err := store.EnsureRegistration(13, reg); err != nil {
@@ -60,7 +69,7 @@ func TestEnsureRegistration(t *testing.T) {
 	if idx != 13 {
 		t.Fatalf("bad: %v", idx)
 	}
-	if len(checks) != 1 {
+	if len(checks) != 2 {
 		t.Fatalf("check: %#v", checks)
 	}
 }