From ef6399928d1584944180b8f0ac6621e550ca8c07 Mon Sep 17 00:00:00 2001
From: Fouad Zaryouh <fouad.zaryouh@gmail.com>
Date: Mon, 3 Feb 2020 04:11:40 -0500
Subject: [PATCH] api: add replace-existing-checks param to the api package
 (#7136)

---
 api/agent.go      | 25 +++++++++++++++
 api/agent_test.go | 82 ++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/api/agent.go b/api/agent.go
index 75999e8e40..9df6ad1df4 100644
--- a/api/agent.go
+++ b/api/agent.go
@@ -159,6 +159,14 @@ type AgentServiceRegistration struct {
 	Namespace         string                          `json:",omitempty" bexpr:"-" hash:"ignore"`
 }
 
+//ServiceRegisterOpts is used to pass extra options to the service register.
+type ServiceRegisterOpts struct {
+	//Missing healthchecks will be deleted from the agent.
+	//Using this parameter allows to idempotently register a service and its checks without
+	//having to manually deregister checks.
+	ReplaceExistingChecks bool
+}
+
 // AgentCheckRegistration is used to register a new check
 type AgentCheckRegistration struct {
 	ID        string `json:",omitempty"`
@@ -554,8 +562,25 @@ func (a *Agent) MembersOpts(opts MembersOpts) ([]*AgentMember, error) {
 // ServiceRegister is used to register a new service with
 // the local agent
 func (a *Agent) ServiceRegister(service *AgentServiceRegistration) error {
+	opts := ServiceRegisterOpts{
+		ReplaceExistingChecks: false,
+	}
+
+	return a.serviceRegister(service, opts)
+}
+
+// ServiceRegister is used to register a new service with
+// the local agent and can be passed additional options.
+func (a *Agent) ServiceRegisterOpts(service *AgentServiceRegistration, opts ServiceRegisterOpts) error {
+	return a.serviceRegister(service, opts)
+}
+
+func (a *Agent) serviceRegister(service *AgentServiceRegistration, opts ServiceRegisterOpts) error {
 	r := a.c.newRequest("PUT", "/v1/agent/service/register")
 	r.obj = service
+	if opts.ReplaceExistingChecks {
+		r.params.Set("replace-existing-checks", "true")
+	}
 	_, resp, err := requireOK(a.c.doRequest(r))
 	if err != nil {
 		return err
diff --git a/api/agent_test.go b/api/agent_test.go
index f9c591a850..085e991fdf 100644
--- a/api/agent_test.go
+++ b/api/agent_test.go
@@ -168,6 +168,86 @@ func TestAPI_AgentMembers(t *testing.T) {
 	}
 }
 
+func TestAPI_AgentServiceAndReplaceChecks(t *testing.T) {
+	t.Parallel()
+	c, s := makeClient(t)
+	defer s.Stop()
+
+	agent := c.Agent()
+	s.WaitForSerfCheck(t)
+
+	reg := &AgentServiceRegistration{
+		Name: "foo",
+		ID:   "foo",
+		Tags: []string{"bar", "baz"},
+		TaggedAddresses: map[string]ServiceAddress{
+			"lan": ServiceAddress{
+				Address: "198.18.0.1",
+				Port:    80,
+			},
+		},
+		Port: 8000,
+		Check: &AgentServiceCheck{
+			TTL: "15s",
+		},
+	}
+
+	regupdate := &AgentServiceRegistration{
+		Name: "foo",
+		ID:   "foo",
+		Tags: []string{"bar", "baz"},
+		TaggedAddresses: map[string]ServiceAddress{
+			"lan": ServiceAddress{
+				Address: "198.18.0.1",
+				Port:    80,
+			},
+		},
+		Port: 9000,
+	}
+
+	if err := agent.ServiceRegister(reg); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if err := agent.ServiceRegisterOpts(regupdate, ServiceRegisterOpts{ReplaceExistingChecks: true}); 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 len(checks) != 0 {
+		t.Fatalf("checks are not removed: %v", checks)
+	}
+
+	state, out, err := agent.AgentHealthServiceByID("foo")
+	require.Nil(t, err)
+	require.NotNil(t, out)
+	require.Equal(t, HealthPassing, state)
+	require.Equal(t, 9000, out.Service.Port)
+
+	state, outs, err := agent.AgentHealthServiceByName("foo")
+	require.Nil(t, err)
+	require.NotNil(t, outs)
+	require.Equal(t, HealthPassing, state)
+	require.Equal(t, 9000, outs[0].Service.Port)
+
+	if err := agent.ServiceDeregister("foo"); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+}
+
 func TestAPI_AgentServices(t *testing.T) {
 	t.Parallel()
 	c, s := makeClient(t)
@@ -231,7 +311,7 @@ func TestAPI_AgentServices(t *testing.T) {
 	require.Nil(t, err)
 	require.NotNil(t, outs)
 	require.Equal(t, HealthCritical, state)
-	require.Equal(t, 8000, out.Service.Port)
+	require.Equal(t, 8000, outs[0].Service.Port)
 
 	if err := agent.ServiceDeregister("foo"); err != nil {
 		t.Fatalf("err: %v", err)