From 197a5a9a9abcca189fd36df3a9fc21791daf3760 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Tue, 13 Jan 2015 19:18:46 -0800
Subject: [PATCH] api: support multiple checks during service/check
 registration

---
 api/agent.go      | 19 ++++++-----
 api/agent_test.go | 81 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 8 deletions(-)

diff --git a/api/agent.go b/api/agent.go
index c31395e189..b71f2c32d7 100644
--- a/api/agent.go
+++ b/api/agent.go
@@ -41,18 +41,20 @@ type AgentMember struct {
 
 // AgentServiceRegistration is used to register a new service
 type AgentServiceRegistration struct {
-	ID    string   `json:",omitempty"`
-	Name  string   `json:",omitempty"`
-	Tags  []string `json:",omitempty"`
-	Port  int      `json:",omitempty"`
-	Check *AgentServiceCheck
+	ID     string   `json:",omitempty"`
+	Name   string   `json:",omitempty"`
+	Tags   []string `json:",omitempty"`
+	Port   int      `json:",omitempty"`
+	Check  *AgentServiceCheck
+	Checks AgentServiceChecks
 }
 
 // AgentCheckRegistration is used to register a new check
 type AgentCheckRegistration struct {
-	ID    string `json:",omitempty"`
-	Name  string `json:",omitempty"`
-	Notes string `json:",omitempty"`
+	ID        string `json:",omitempty"`
+	Name      string `json:",omitempty"`
+	Notes     string `json:",omitempty"`
+	ServiceID string `json:",omitempty"`
 	AgentServiceCheck
 }
 
@@ -63,6 +65,7 @@ type AgentServiceCheck struct {
 	Interval string `json:",omitempty"`
 	TTL      string `json:",omitempty"`
 }
+type AgentServiceChecks []*AgentServiceCheck
 
 // Agent can be used to query the Agent endpoints
 type Agent struct {
diff --git a/api/agent_test.go b/api/agent_test.go
index 06f8ca12e1..627256c230 100644
--- a/api/agent_test.go
+++ b/api/agent_test.go
@@ -76,6 +76,49 @@ func TestAgent_Services(t *testing.T) {
 	}
 }
 
+func TestAgent_Services_MultipleChecks(t *testing.T) {
+	c, s := makeClient(t)
+	defer s.stop()
+
+	agent := c.Agent()
+
+	reg := &AgentServiceRegistration{
+		Name: "foo",
+		Tags: []string{"bar", "baz"},
+		Port: 8000,
+		Checks: AgentServiceChecks{
+			&AgentServiceCheck{
+				TTL: "15s",
+			},
+			&AgentServiceCheck{
+				TTL: "30s",
+			},
+		},
+	}
+	if err := agent.ServiceRegister(reg); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	services, err := agent.Services()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if _, ok := services["foo"]; !ok {
+		t.Fatalf("missing service: %v", services)
+	}
+
+	checks, err := agent.Checks()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	if _, ok := checks["service:foo:1"]; !ok {
+		t.Fatalf("missing check: %v", checks)
+	}
+	if _, ok := checks["service:foo:2"]; !ok {
+		t.Fatalf("missing check: %v", checks)
+	}
+}
+
 func TestAgent_SetTTLStatus(t *testing.T) {
 	c, s := makeClient(t)
 	defer s.stop()
@@ -143,6 +186,44 @@ func TestAgent_Checks(t *testing.T) {
 	}
 }
 
+func TestAgent_Checks_serviceBound(t *testing.T) {
+	c, s := makeClient(t)
+	defer s.stop()
+
+	agent := c.Agent()
+
+	// First register a service
+	serviceReg := &AgentServiceRegistration{
+		Name: "redis",
+	}
+	if err := agent.ServiceRegister(serviceReg); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Register a check bound to the service
+	reg := &AgentCheckRegistration{
+		Name:      "redischeck",
+		ServiceID: "redis",
+	}
+	reg.TTL = "15s"
+	if err := agent.CheckRegister(reg); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	checks, err := agent.Checks()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	check, ok := checks["redischeck"]
+	if !ok {
+		t.Fatalf("missing check: %v", checks)
+	}
+	if check.ServiceID != "redis" {
+		t.Fatalf("missing service association for check: %v", check)
+	}
+}
+
 func TestAgent_Join(t *testing.T) {
 	c, s := makeClient(t)
 	defer s.stop()