From d0a28482ff955609091545213e38cdacfde79a93 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Thu, 7 May 2015 15:30:01 -0700
Subject: [PATCH] agent: use service ID field to determine associated health
 checks during deregister

---
 command/agent/agent.go      |  6 ++----
 command/agent/agent_test.go | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index e52ca60f3..eb19875dd 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -10,7 +10,6 @@ import (
 	"path/filepath"
 	"regexp"
 	"strconv"
-	"strings"
 	"sync"
 
 	"github.com/hashicorp/consul/consul"
@@ -703,9 +702,8 @@ func (a *Agent) RemoveService(serviceID string, persist bool) error {
 	}
 
 	// Deregister any associated health checks
-	for checkID, _ := range a.state.Checks() {
-		prefix := "service:" + serviceID
-		if checkID != prefix && !strings.HasPrefix(checkID, prefix+":") {
+	for checkID, health := range a.state.Checks() {
+		if health.ServiceID != serviceID {
 			continue
 		}
 		if err := a.RemoveCheck(checkID, persist); err != nil {
diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go
index dd3ed1a75..af1f206e2 100644
--- a/command/agent/agent_test.go
+++ b/command/agent/agent_test.go
@@ -265,12 +265,27 @@ func TestAgent_RemoveService(t *testing.T) {
 			t.Fatalf("err: %v", err)
 		}
 
+		// Add a check after the fact with a specific check ID
+		check := &CheckDefinition{
+			ID:        "check2",
+			Name:      "check2",
+			ServiceID: "memcache",
+			CheckType: CheckType{TTL: time.Minute},
+		}
+		hc := check.HealthCheck("node1")
+		if err := agent.AddCheck(hc, &check.CheckType, false, ""); err != nil {
+			t.Fatalf("err: %s", err)
+		}
+
 		if err := agent.RemoveService("memcache", false); err != nil {
 			t.Fatalf("err: %s", err)
 		}
 		if _, ok := agent.state.Checks()["service:memcache"]; ok {
 			t.Fatalf("have memcache check")
 		}
+		if _, ok := agent.state.Checks()["check2"]; ok {
+			t.Fatalf("have check2 check")
+		}
 	}
 
 	// Removing a service with multiple checks works