From 1f02461f8c8e5b6b1cc73eb79884b64b84c66bdf Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Wed, 8 Jan 2014 12:19:11 -0800
Subject: [PATCH] Support getting node checks in the snapshot

---
 consul/state_store.go      |  5 +++++
 consul/state_store_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/consul/state_store.go b/consul/state_store.go
index 565ad88629..26745007ee 100644
--- a/consul/state_store.go
+++ b/consul/state_store.go
@@ -518,3 +518,8 @@ func (s *StateSnapshot) Nodes() structs.Nodes {
 func (s *StateSnapshot) NodeServices(name string) *structs.NodeServices {
 	return s.store.parseNodeServices(s.tx, name)
 }
+
+// NodeChecks is used to return all the checks of a given node
+func (s *StateSnapshot) NodeChecks(node string) structs.HealthChecks {
+	return parseHealthChecks(s.store.checkTable.GetTxn(s.tx, "id", node))
+}
diff --git a/consul/state_store_test.go b/consul/state_store_test.go
index ac1dda7441..fc2279acf1 100644
--- a/consul/state_store_test.go
+++ b/consul/state_store_test.go
@@ -479,6 +479,17 @@ func TestStoreSnapshot(t *testing.T) {
 		t.Fatalf("err: %v")
 	}
 
+	check := &structs.HealthCheck{
+		Node:      "foo",
+		CheckID:   "db",
+		Name:      "Can connect",
+		Status:    structs.HealthPassing,
+		ServiceID: "db",
+	}
+	if err := store.EnsureCheck(check); err != nil {
+		t.Fatalf("err: %v")
+	}
+
 	// Take a snapshot
 	snap, err := store.Snapshot()
 	if err != nil {
@@ -506,6 +517,15 @@ func TestStoreSnapshot(t *testing.T) {
 		t.Fatalf("bad: %v", services)
 	}
 
+	// Ensure we get the checks
+	checks := snap.NodeChecks("foo")
+	if len(checks) != 1 {
+		t.Fatalf("bad: %v", checks)
+	}
+	if !reflect.DeepEqual(checks[0], check) {
+		t.Fatalf("bad: %v", checks[0])
+	}
+
 	// Make some changes!
 	if err := store.EnsureService("foo", "db", "db", "slave", 8000); err != nil {
 		t.Fatalf("err: %v", err)
@@ -516,6 +536,16 @@ func TestStoreSnapshot(t *testing.T) {
 	if err := store.EnsureNode(structs.Node{"baz", "127.0.0.3"}); err != nil {
 		t.Fatalf("err: %v", err)
 	}
+	checkAfter := &structs.HealthCheck{
+		Node:      "foo",
+		CheckID:   "db",
+		Name:      "Can connect",
+		Status:    structs.HealthCritical,
+		ServiceID: "db",
+	}
+	if err := store.EnsureCheck(checkAfter); err != nil {
+		t.Fatalf("err: %v")
+	}
 
 	// Check snapshot has old values
 	nodes = snap.Nodes()
@@ -536,6 +566,14 @@ func TestStoreSnapshot(t *testing.T) {
 	if services.Services["db"].Tag != "slave" {
 		t.Fatalf("bad: %v", services)
 	}
+
+	checks = snap.NodeChecks("foo")
+	if len(checks) != 1 {
+		t.Fatalf("bad: %v", checks)
+	}
+	if !reflect.DeepEqual(checks[0], check) {
+		t.Fatalf("bad: %v", checks[0])
+	}
 }
 
 func TestEnsureCheck(t *testing.T) {