From e5f5fc9301c70e962652059695a3be28d57c4e0d Mon Sep 17 00:00:00 2001
From: Tim Gross <tgross@hashicorp.com>
Date: Mon, 25 Sep 2023 11:25:02 -0400
Subject: [PATCH] api: add `CheckRegisterOpts` method to Agent API (#18943)

Ongoing work to support Nomad Workload Identity for authenticating with Consul
will mean that Nomad's service registration sync with Consul will want to use
Consul tokens scoped to individual workloads for registering services and
checks. The `CheckRegister` method in the API doesn't have an option to pass the
token in, which prevent us from sharing the same Consul connection for all
workloads. Add a `CheckRegisterOpts` to match the behavior of
`ServiceRegisterOpts`.
---
 .changelog/18943.txt |  3 +++
 api/agent.go         |  7 +++++++
 api/agent_test.go    | 15 ++++++++++++++-
 3 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 .changelog/18943.txt

diff --git a/.changelog/18943.txt b/.changelog/18943.txt
new file mode 100644
index 0000000000..adb027dee9
--- /dev/null
+++ b/.changelog/18943.txt
@@ -0,0 +1,3 @@
+```release-note:improvement
+api: added `CheckRegisterOpts` to Agent API
+```
diff --git a/api/agent.go b/api/agent.go
index 8c2bf6b455..24e2c50d64 100644
--- a/api/agent.go
+++ b/api/agent.go
@@ -998,7 +998,14 @@ func (a *Agent) UpdateTTLOpts(checkID, output, status string, q *QueryOptions) e
 // CheckRegister is used to register a new check with
 // the local agent
 func (a *Agent) CheckRegister(check *AgentCheckRegistration) error {
+	return a.CheckRegisterOpts(check, nil)
+}
+
+// CheckRegisterOpts is used to register a new check with
+// the local agent using query options
+func (a *Agent) CheckRegisterOpts(check *AgentCheckRegistration, q *QueryOptions) error {
 	r := a.c.newRequest("PUT", "/v1/agent/check/register")
+	r.setQueryOptions(q)
 	r.obj = check
 	_, resp, err := a.c.doRequest(r)
 	if err != nil {
diff --git a/api/agent_test.go b/api/agent_test.go
index 2c359fe86e..511d51607d 100644
--- a/api/agent_test.go
+++ b/api/agent_test.go
@@ -1073,7 +1073,7 @@ func TestAPI_AgentChecks(t *testing.T) {
 		Name: "foo",
 	}
 	reg.TTL = "15s"
-	if err := agent.CheckRegister(reg); err != nil {
+	if err := agent.CheckRegisterOpts(reg, nil); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
@@ -1097,6 +1097,19 @@ func TestAPI_AgentChecks(t *testing.T) {
 	}
 }
 
+func TestAgent_AgentChecksRegisterOpts_WithContextTimeout(t *testing.T) {
+	c, err := NewClient(DefaultConfig())
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
+	t.Cleanup(cancel)
+
+	opts := &QueryOptions{}
+	opts = opts.WithContext(ctx)
+	err = c.Agent().CheckRegisterOpts(&AgentCheckRegistration{}, opts)
+	require.True(t, errors.Is(err, context.DeadlineExceeded), "expected timeout")
+}
+
 func TestAPI_AgentChecksWithFilterOpts(t *testing.T) {
 	t.Parallel()
 	c, s := makeClient(t)