From 377c3a9b0b69a650deab81531d3fd5ac2fc01b36 Mon Sep 17 00:00:00 2001
From: FFMMM <FFMMM@users.noreply.github.com>
Date: Mon, 20 Sep 2021 14:04:13 -0700
Subject: [PATCH] add StatusError to api package (#11054)

* add require http codes in api and use in operator_autopilot health check

* add StatusError type in api package

Signed-off-by: FFMMM <FFMMM@users.noreply.github.com>
---
 api/api.go                     | 31 ++++++++++++++++++++++++++-----
 api/operator_autopilot.go      | 13 ++++---------
 api/operator_autopilot_test.go | 12 ++++++++++++
 3 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/api/api.go b/api/api.go
index 9d84c286df..3063618744 100644
--- a/api/api.go
+++ b/api/api.go
@@ -82,6 +82,15 @@ const (
 	HTTPPartitionEnvName = "CONSUL_PARTITION"
 )
 
+type StatusError struct {
+	Code int
+	Body string
+}
+
+func (e StatusError) Error() string {
+	return fmt.Sprintf("Unexpected response code: %d (%s)", e.Code, e.Body)
+}
+
 // QueryOptions are used to parameterize a query
 type QueryOptions struct {
 	// Namespace overrides the `default` namespace
@@ -972,7 +981,7 @@ func (c *Client) doRequest(r *request) (time.Duration, *http.Response, error) {
 func (c *Client) query(endpoint string, out interface{}, q *QueryOptions) (*QueryMeta, error) {
 	r := c.newRequest("GET", endpoint)
 	r.setQueryOptions(q)
-	rtt, resp, err := c.doRequest(r)
+	rtt, resp, err := requireOK(c.doRequest(r))
 	if err != nil {
 		return nil, err
 	}
@@ -1090,16 +1099,28 @@ func encodeBody(obj interface{}) (io.Reader, error) {
 
 // requireOK is used to wrap doRequest and check for a 200
 func requireOK(d time.Duration, resp *http.Response, e error) (time.Duration, *http.Response, error) {
+	return requireHttpCodes(d, resp, e, 200)
+}
+
+// requireHttpCodes checks for the "allowable" http codes for a response
+func requireHttpCodes(d time.Duration, resp *http.Response, e error, httpCodes ...int) (time.Duration, *http.Response, error) {
 	if e != nil {
 		if resp != nil {
 			closeResponseBody(resp)
 		}
 		return d, nil, e
 	}
-	if resp.StatusCode != 200 {
-		return d, nil, generateUnexpectedResponseCodeError(resp)
+
+	// if there is an http code that we require, return w no error
+	for _, httpCode := range httpCodes {
+		if resp.StatusCode == httpCode {
+			return d, resp, nil
+		}
 	}
-	return d, resp, nil
+
+	// if we reached here, then none of the http codes in resp matched any that we expected
+	// so err out
+	return d, nil, generateUnexpectedResponseCodeError(resp)
 }
 
 // closeResponseBody reads resp.Body until EOF, and then closes it. The read
@@ -1127,7 +1148,7 @@ func generateUnexpectedResponseCodeError(resp *http.Response) error {
 	closeResponseBody(resp)
 
 	trimmed := strings.TrimSpace(string(buf.Bytes()))
-	return fmt.Errorf("Unexpected response code: %d (%s)", resp.StatusCode, trimmed)
+	return StatusError{Code: resp.StatusCode, Body: trimmed}
 }
 
 func requireNotFoundOrOK(d time.Duration, resp *http.Response, e error) (bool, time.Duration, *http.Response, error) {
diff --git a/api/operator_autopilot.go b/api/operator_autopilot.go
index d713ee1cc3..99b02faab8 100644
--- a/api/operator_autopilot.go
+++ b/api/operator_autopilot.go
@@ -352,9 +352,10 @@ func (op *Operator) AutopilotServerHealth(q *QueryOptions) (*OperatorHealthReply
 	r := op.c.newRequest("GET", "/v1/operator/autopilot/health")
 	r.setQueryOptions(q)
 
-	// we cannot just use requireOK because this endpoint might use a 429 status to indicate
-	// that unhealthiness
-	_, resp, err := op.c.doRequest(r)
+	// we use 429 status to indicate unhealthiness
+	d, resp, err := op.c.doRequest(r)
+	_, resp, err = requireHttpCodes(d, resp, err, 200, 429)
+
 	if err != nil {
 		if resp != nil {
 			closeResponseBody(resp)
@@ -362,12 +363,6 @@ func (op *Operator) AutopilotServerHealth(q *QueryOptions) (*OperatorHealthReply
 		return nil, err
 	}
 
-	// these are the only 2 status codes that would indicate that we should
-	// expect the body to contain the right format.
-	if resp.StatusCode != 200 && resp.StatusCode != 429 {
-		return nil, generateUnexpectedResponseCodeError(resp)
-	}
-
 	defer closeResponseBody(resp)
 
 	var out OperatorHealthReply
diff --git a/api/operator_autopilot_test.go b/api/operator_autopilot_test.go
index 720ae29413..a0850110ff 100644
--- a/api/operator_autopilot_test.go
+++ b/api/operator_autopilot_test.go
@@ -1,6 +1,7 @@
 package api
 
 import (
+	"errors"
 	"testing"
 	"time"
 
@@ -181,4 +182,15 @@ func TestAPI_OperatorAutopilotServerHealth_429(t *testing.T) {
 	out, err := client.Operator().AutopilotServerHealth(nil)
 	require.NoError(t, err)
 	require.Equal(t, &reply, out)
+
+	mapi.withReply("GET", "/v1/operator/autopilot/health", nil, 500, nil).Once()
+	_, err = client.Operator().AutopilotServerHealth(nil)
+
+	var statusE StatusError
+	if errors.As(err, &statusE) {
+		require.Equal(t, 500, statusE.Code)
+	} else {
+		t.Error("Failed to unwrap error as StatusError")
+	}
+
 }