From a24f6e3d4ded207e28ce639459fd88bf102da8d5 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Wed, 7 Jan 2015 22:26:40 -0800
Subject: [PATCH] agent: test service and check unloading

---
 command/agent/agent.go      | 16 +++++--
 command/agent/agent_test.go | 91 +++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 5 deletions(-)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index 6b4d4a672f..c34063a861 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -622,7 +622,12 @@ func (a *Agent) RemoveService(serviceID string, persist bool) error {
 
 	// Deregister any associated health checks
 	checkID := fmt.Sprintf("service:%s", serviceID)
-	return a.RemoveCheck(checkID, persist)
+	if err := a.RemoveCheck(checkID, persist); err != nil {
+		return err
+	}
+
+	log.Printf("[DEBUG] agent: removed service %q", serviceID)
+	return nil
 }
 
 // AddCheck is used to add a health check to the agent.
@@ -710,6 +715,7 @@ func (a *Agent) RemoveCheck(checkID string, persist bool) error {
 	if persist {
 		return a.purgeCheck(checkID)
 	}
+	log.Printf("[DEBUG] agent: removed check %q", checkID)
 	return nil
 }
 
@@ -850,11 +856,11 @@ func (a *Agent) loadServices(conf *Config) error {
 		if _, ok := a.state.services[svc.ID]; ok {
 			// Purge previously persisted service. This allows config to be
 			// preferred over services persisted from the API.
-			a.logger.Printf("[DEBUG] Service %q exists, not restoring from %q",
+			a.logger.Printf("[DEBUG] agent: service %q exists, not restoring from %q",
 				svc.ID, filePath)
 			return a.purgeService(svc.ID)
 		} else {
-			a.logger.Printf("[DEBUG] Restored service definition %q from %q",
+			a.logger.Printf("[DEBUG] agent: restored service definition %q from %q",
 				svc.ID, filePath)
 			return a.AddService(svc, nil, false)
 		}
@@ -921,7 +927,7 @@ func (a *Agent) loadChecks(conf *Config) error {
 		if _, ok := a.state.checks[p.Check.CheckID]; ok {
 			// Purge previously persisted check. This allows config to be
 			// preferred over persisted checks from the API.
-			a.logger.Printf("[DEBUG] Check %q exists, not restoring from %q",
+			a.logger.Printf("[DEBUG] agent: check %q exists, not restoring from %q",
 				p.Check.CheckID, filePath)
 			return a.purgeCheck(p.Check.CheckID)
 		} else {
@@ -929,7 +935,7 @@ func (a *Agent) loadChecks(conf *Config) error {
 			// services into the active pool
 			p.Check.Status = structs.HealthCritical
 
-			a.logger.Printf("[DEBUG] Restored health check %q from %q",
+			a.logger.Printf("[DEBUG] agent: restored health check %q from %q",
 				p.Check.CheckID, filePath)
 			return a.AddCheck(p.Check, p.ChkType, false)
 		}
diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go
index e2a7ef70e0..0add36702c 100644
--- a/command/agent/agent_test.go
+++ b/command/agent/agent_test.go
@@ -690,3 +690,94 @@ func TestAgent_PurgeCheckOnDuplicate(t *testing.T) {
 		t.Fatalf("bad: %#v", result)
 	}
 }
+
+func TestAgent_unloadChecks(t *testing.T) {
+	config := nextConfig()
+	dir, agent := makeAgent(t, config)
+	defer os.RemoveAll(dir)
+	defer agent.Shutdown()
+
+	check1 := &structs.HealthCheck{
+		Node:        config.NodeName,
+		CheckID:     "service:redis1",
+		Name:        "redischeck",
+		Status:      structs.HealthPassing,
+		ServiceID:   "redis",
+		ServiceName: "redis",
+	}
+
+	// Register the check
+	if err := agent.AddCheck(check1, nil, false); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	found := false
+	for check, _ := range agent.state.Checks() {
+		if check == check1.CheckID {
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Fatalf("check should have been registered")
+	}
+
+	// Unload all of the checks
+	if err := agent.unloadChecks(); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Make sure it was unloaded
+	for check, _ := range agent.state.Checks() {
+		if check == check1.CheckID {
+			t.Fatalf("should have unloaded checks")
+		}
+	}
+}
+
+func TestAgent_unloadServices(t *testing.T) {
+	config := nextConfig()
+	dir, agent := makeAgent(t, config)
+	defer os.RemoveAll(dir)
+	defer agent.Shutdown()
+
+	svc := &structs.NodeService{
+		ID:      "redis",
+		Service: "redis",
+		Tags:    []string{"foo"},
+		Port:    8000,
+	}
+
+	// Register the service
+	if err := agent.AddService(svc, nil, false); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	found := false
+	for id, _ := range agent.state.Services() {
+		if id == svc.ID {
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Fatalf("should have registered service")
+	}
+
+	// Unload all services
+	if err := agent.unloadServices(); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Make sure it was unloaded and the consul service remains
+	found = false
+	for id, _ := range agent.state.Services() {
+		if id == svc.ID {
+			t.Fatalf("should have unloaded services")
+		}
+		if id == consul.ConsulServiceID {
+			found = true
+		}
+	}
+	if !found {
+		t.Fatalf("consul service should not be removed")
+	}
+}