Merge branch 'main' into APICLI-327/token-validation

CodeBetterEveryDay · Dec 2, 2020 · fcbbb1e · fcbbb1e
2 parents 019045d + a11a281
commit fcbbb1e
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 11 deletions.
diff --git a/godo.go b/godo.go
@@ -80,6 +80,9 @@ type Client struct {
 
 	// Optional function called after every successful request made to the DO APIs
 	onRequestCompleted RequestCompletionCallback
+
+	// Optional extra HTTP headers to set on every request to the API.
+	headers map[string]string
 }
 
 // RequestCompletionCallback defines the type of the request callback function
@@ -217,6 +220,8 @@ func NewClient(httpClient *http.Client) *Client {
 	c.VPCs = &VPCsServiceOp{client: c}
 	c.OneClick = &OneClickServiceOp{client: c}
 
+	c.headers = make(map[string]string)
+
 	return c
 }
 
@@ -256,6 +261,17 @@ func SetUserAgent(ua string) ClientOpt {
 	}
 }
 
+// SetRequestHeaders sets optional HTTP headers on the client that are
+// sent on each HTTP request.
+func SetRequestHeaders(headers map[string]string) ClientOpt {
+	return func(c *Client) error {
+		for k, v := range headers {
+			c.headers[k] = v
+		}
+		return nil
+	}
+}
+
 // NewRequest creates an API request. A relative URL can be provided in urlStr, which will be resolved to the
 // BaseURL of the Client. Relative URLS should always be specified without a preceding slash. If specified, the
 // value pointed to by body is JSON encoded and included in as the request body.
@@ -265,22 +281,37 @@ func (c *Client) NewRequest(ctx context.Context, method, urlStr string, body int
 		return nil, err
 	}
 
-	buf := new(bytes.Buffer)
-	if body != nil {
-		err = json.NewEncoder(buf).Encode(body)
+	var req *http.Request
+	switch method {
+	case http.MethodGet, http.MethodHead, http.MethodOptions:
+		req, err = http.NewRequest(method, u.String(), nil)
 		if err != nil {
 			return nil, err
 		}
+
+	default:
+		buf := new(bytes.Buffer)
+		if body != nil {
+			err = json.NewEncoder(buf).Encode(body)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		req, err = http.NewRequest(method, u.String(), buf)
+		if err != nil {
+			return nil, err
+		}
+		req.Header.Set("Content-Type", mediaType)
 	}
 
-	req, err := http.NewRequest(method, u.String(), buf)
-	if err != nil {
-		return nil, err
+	for k, v := range c.headers {
+		req.Header.Add(k, v)
 	}
 
-	req.Header.Add("Content-Type", mediaType)
-	req.Header.Add("Accept", mediaType)
-	req.Header.Add("User-Agent", c.UserAgent)
+	req.Header.Set("Accept", mediaType)
+	req.Header.Set("User-Agent", c.UserAgent)
+
 	return req, nil
 }
 

diff --git a/godo_test.go b/godo_test.go
@@ -175,7 +175,7 @@ func TestNewRequest(t *testing.T) {
 		`{"name":"l","region":"","size":"","image":0,`+
 			`"ssh_keys":null,"backups":false,"ipv6":false,`+
 			`"private_networking":false,"monitoring":false,"tags":null}`+"\n"
-	req, _ := c.NewRequest(ctx, http.MethodGet, inURL, inBody)
+	req, _ := c.NewRequest(ctx, http.MethodPost, inURL, inBody)
 
 	// test relative URL was expanded
 	if req.URL.String() != outURL {
@@ -195,6 +195,29 @@ func TestNewRequest(t *testing.T) {
 	}
 }
 
+func TestNewRequest_get(t *testing.T) {
+	c := NewClient(nil)
+
+	inURL, outURL := "/foo", defaultBaseURL+"foo"
+	req, _ := c.NewRequest(ctx, http.MethodGet, inURL, nil)
+
+	// test relative URL was expanded
+	if req.URL.String() != outURL {
+		t.Errorf("NewRequest(%v) URL = %v, expected %v", inURL, req.URL, outURL)
+	}
+
+	// test the content-type header is not set
+	if contentType := req.Header.Get("Content-Type"); contentType != "" {
+		t.Errorf("NewRequest() Content-Type = %v, expected empty string", contentType)
+	}
+
+	// test default user-agent is attached to the request
+	userAgent := req.Header.Get("User-Agent")
+	if c.UserAgent != userAgent {
+		t.Errorf("NewRequest() User-Agent = %v, expected %v", userAgent, c.UserAgent)
+	}
+}
+
 func TestNewRequest_withUserData(t *testing.T) {
 	c := NewClient(nil)
 
@@ -203,7 +226,7 @@ func TestNewRequest_withUserData(t *testing.T) {
 		`{"name":"l","region":"","size":"","image":0,`+
 			`"ssh_keys":null,"backups":false,"ipv6":false,`+
 			`"private_networking":false,"monitoring":false,"user_data":"u","tags":null}`+"\n"
-	req, _ := c.NewRequest(ctx, http.MethodGet, inURL, inBody)
+	req, _ := c.NewRequest(ctx, http.MethodPost, inURL, inBody)
 
 	// test relative URL was expanded
 	if req.URL.String() != outURL {
@@ -245,6 +268,28 @@ func TestNewRequest_withCustomUserAgent(t *testing.T) {
 	}
 }
 
+func TestNewRequest_withCustomHeaders(t *testing.T) {
+	expectedIdentity := "identity"
+	expectedCustom := "x_test_header"
+
+	c, err := New(nil, SetRequestHeaders(map[string]string{
+		"Accept-Encoding": expectedIdentity,
+		"X-Test-Header":   expectedCustom,
+	}))
+	if err != nil {
+		t.Fatalf("New() unexpected error: %v", err)
+	}
+
+	req, _ := c.NewRequest(ctx, http.MethodGet, "/foo", nil)
+
+	if got := req.Header.Get("Accept-Encoding"); got != expectedIdentity {
+		t.Errorf("New() Custom Accept Encoding Header = %s; expected %s", got, expectedIdentity)
+	}
+	if got := req.Header.Get("X-Test-Header"); got != expectedCustom {
+		t.Errorf("New() Custom Accept Encoding Header = %s; expected %s", got, expectedCustom)
+	}
+}
+
 func TestDo(t *testing.T) {
 	setup()
 	defer teardown()