From 35f5a65fb71f7539d10c7d95811d50c67387ce62 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Tue, 28 Apr 2015 13:06:02 -0700
Subject: [PATCH] agent: more tests

---
 command/agent/agent_endpoint_test.go | 14 ++++++++++--
 command/agent/local.go               | 22 ++++++++++++++++--
 command/agent/local_test.go          | 34 ++++++++++++++++++++++++++--
 3 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/command/agent/agent_endpoint_test.go b/command/agent/agent_endpoint_test.go
index 0387f38901..5c95c1a3e7 100644
--- a/command/agent/agent_endpoint_test.go
+++ b/command/agent/agent_endpoint_test.go
@@ -252,7 +252,7 @@ func TestHTTPAgentRegisterCheck(t *testing.T) {
 	defer srv.agent.Shutdown()
 
 	// Register node
-	req, err := http.NewRequest("GET", "/v1/agent/check/register", nil)
+	req, err := http.NewRequest("GET", "/v1/agent/check/register?token=abc123", nil)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -280,6 +280,11 @@ func TestHTTPAgentRegisterCheck(t *testing.T) {
 	if _, ok := srv.agent.checkTTLs["test"]; !ok {
 		t.Fatalf("missing test check ttl")
 	}
+
+	// Ensure the token was configured
+	if token := srv.agent.state.CheckToken("test"); token == "" {
+		t.Fatalf("missing token")
+	}
 }
 
 func TestHTTPAgentDeregisterCheck(t *testing.T) {
@@ -419,7 +424,7 @@ func TestHTTPAgentRegisterService(t *testing.T) {
 	defer srv.agent.Shutdown()
 
 	// Register node
-	req, err := http.NewRequest("GET", "/v1/agent/service/register", nil)
+	req, err := http.NewRequest("GET", "/v1/agent/service/register?token=abc123", nil)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -463,6 +468,11 @@ func TestHTTPAgentRegisterService(t *testing.T) {
 	if len(srv.agent.checkTTLs) != 3 {
 		t.Fatalf("missing test check ttls: %v", srv.agent.checkTTLs)
 	}
+
+	// Ensure the token was configured
+	if token := srv.agent.state.ServiceToken("test"); token == "" {
+		t.Fatalf("missing token")
+	}
 }
 
 func TestHTTPAgentDeregisterService(t *testing.T) {
diff --git a/command/agent/local.go b/command/agent/local.go
index a0c200967e..83fefbf885 100644
--- a/command/agent/local.go
+++ b/command/agent/local.go
@@ -126,9 +126,18 @@ func (l *localState) isPaused() bool {
 // SetServiceToken configures the provided token for the service ID.
 // The token will be used to perform service registration operations.
 func (l *localState) SetServiceToken(id, token string) {
+	if token != "" {
+		l.Lock()
+		defer l.Unlock()
+		l.serviceTokens[id] = token
+	}
+}
+
+// RemoveServiceToken is used to remove a configured service token.
+func (l *localState) RemoveServiceToken(id string) {
 	l.Lock()
 	defer l.Unlock()
-	l.serviceTokens[id] = token
+	delete(l.serviceTokens, id)
 }
 
 // ServiceToken returns the configured ACL token for the given
@@ -192,9 +201,18 @@ func (l *localState) Services() map[string]*structs.NodeService {
 // SetCheckToken is used to configure an ACL token for a specific
 // health check. The token is used during check registration operations.
 func (l *localState) SetCheckToken(id, token string) {
+	if token != "" {
+		l.Lock()
+		defer l.Unlock()
+		l.checkTokens[id] = token
+	}
+}
+
+// RemoveCheckToken is used to remove a configured check token.
+func (l *localState) RemoveCheckToken(id string) {
 	l.Lock()
 	defer l.Unlock()
-	l.checkTokens[id] = token
+	delete(l.checkTokens, id)
 }
 
 // CheckToken is used to return the configured health check token, or
diff --git a/command/agent/local_test.go b/command/agent/local_test.go
index e66977e5a8..84c7c2d26a 100644
--- a/command/agent/local_test.go
+++ b/command/agent/local_test.go
@@ -616,21 +616,51 @@ func TestAgentAntiEntropy_deleteCheck_fails(t *testing.T) {
 }
 
 func TestAgent_serviceTokens(t *testing.T) {
+	config := nextConfig()
+	config.ACLToken = "default"
 	l := new(localState)
-	l.Init(nil, nil)
+	l.Init(config, nil)
+
+	// Returns default when no token is set
+	if token := l.ServiceToken("redis"); token != "default" {
+		t.Fatalf("bad: %s", token)
+	}
+
+	// Returns configured token
 	l.SetServiceToken("redis", "abc123")
 	if token := l.ServiceToken("redis"); token != "abc123" {
 		t.Fatalf("bad: %s", token)
 	}
+
+	// Removes token
+	l.RemoveServiceToken("redis")
+	if token := l.ServiceToken("redis"); token != "default" {
+		t.Fatalf("bad: %s", token)
+	}
 }
 
 func TestAgent_checkTokens(t *testing.T) {
+	config := nextConfig()
+	config.ACLToken = "default"
 	l := new(localState)
-	l.Init(nil, nil)
+	l.Init(config, nil)
+
+	// Returns default when no token is set
+	if token := l.CheckToken("mem"); token != "default" {
+		t.Fatalf("bad: %s", token)
+	}
+
+	// Returns configured token
 	l.SetCheckToken("mem", "abc123")
 	if token := l.CheckToken("mem"); token != "abc123" {
 		t.Fatalf("bad: %s", token)
 	}
+
+	// Removes token
+	l.RemoveCheckToken("mem")
+	if token := l.CheckToken("mem"); token != "default" {
+		t.Fatalf("bad: %s", token)
+	}
 }
 
 var testRegisterRules = `