From 76cb81bcd4ca31ef15731291f9935cf7a55640e1 Mon Sep 17 00:00:00 2001
From: Tim Heckman <t@heckman.io>
Date: Sat, 27 Feb 2021 11:24:43 -0800
Subject: [PATCH] Update extensions.go to accept a context.Context

Updates #267
---
 extension.go | 53 +++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 46 insertions(+), 7 deletions(-)

diff --git a/extension.go b/extension.go
index d4335779..9afb7eb8 100644
--- a/extension.go
+++ b/extension.go
@@ -8,6 +8,8 @@ import (
 	"github.com/google/go-querystring/query"
 )
 
+// Extension represents a single PagerDuty extension. These are addtional
+// features to be used as part of the incident management process.
 type Extension struct {
 	APIObject
 	Name             string      `json:"name"`
@@ -17,11 +19,14 @@ type Extension struct {
 	Config           interface{} `json:"config"`
 }
 
+// ListExtensionResponse represents the single response from the PagerDuty API
+// when listing extensions.
 type ListExtensionResponse struct {
 	APIListObject
 	Extensions []Extension `json:"extensions"`
 }
 
+// ListExtensionOptions are the options to use when listing extensions.
 type ListExtensionOptions struct {
 	APIListObject
 	ExtensionObjectID string `url:"extension_object_id,omitempty"`
@@ -29,39 +34,73 @@ type ListExtensionOptions struct {
 	Query             string `url:"query,omitempty"`
 }
 
+// ListExtensions lists the extensions from the API. It's recommended to use
+// ListExtensionsWithContext instead.
 func (c *Client) ListExtensions(o ListExtensionOptions) (*ListExtensionResponse, error) {
+	return c.ListExtensionsWithContext(context.Background(), o)
+}
+
+// ListExtensionsWithContext lists the extensions from the API.
+func (c *Client) ListExtensionsWithContext(ctx context.Context, o ListExtensionOptions) (*ListExtensionResponse, error) {
 	v, err := query.Values(o)
 	if err != nil {
 		return nil, err
 	}
 
-	resp, err := c.get(context.TODO(), "/extensions?"+v.Encode())
+	resp, err := c.get(ctx, "/extensions?"+v.Encode())
 	if err != nil {
 		return nil, err
 	}
 
 	var result ListExtensionResponse
+	if err = c.decodeJSON(resp, &result); err != nil {
+		return nil, err
+	}
 
-	return &result, c.decodeJSON(resp, &result)
+	return &result, nil
 }
 
+// CreateExtension creates a single extension.
 func (c *Client) CreateExtension(e *Extension) (*Extension, error) {
-	resp, err := c.post(context.TODO(), "/extensions", e, nil)
+	return c.CreateExtensionWithContext(context.Background(), e)
+}
+
+// CreateExtensionWithContext creates a single extension.
+func (c *Client) CreateExtensionWithContext(ctx context.Context, e *Extension) (*Extension, error) {
+	resp, err := c.post(ctx, "/extensions", e, nil)
 	return getExtensionFromResponse(c, resp, err)
 }
 
+// DeleteExtension deletes an extension by its ID.
 func (c *Client) DeleteExtension(id string) error {
-	_, err := c.delete(context.TODO(), "/extensions/"+id)
+	return c.DeleteExtensionWithContext(context.Background(), id)
+}
+
+// DeleteExtensionWithContext deletes an extension by its ID.
+func (c *Client) DeleteExtensionWithContext(ctx context.Context, id string) error {
+	_, err := c.delete(ctx, "/extensions/"+id)
 	return err
 }
 
+// GetExtension gets an extension by its ID.
 func (c *Client) GetExtension(id string) (*Extension, error) {
-	resp, err := c.get(context.TODO(), "/extensions/"+id)
+	return c.GetExtensionWithContext(context.Background(), id)
+}
+
+// GetExtensionWithContext gets an extension by its ID.
+func (c *Client) GetExtensionWithContext(ctx context.Context, id string) (*Extension, error) {
+	resp, err := c.get(ctx, "/extensions/"+id)
 	return getExtensionFromResponse(c, resp, err)
 }
 
+// UpdateExtension updates an extension by its ID.
 func (c *Client) UpdateExtension(id string, e *Extension) (*Extension, error) {
-	resp, err := c.put(context.TODO(), "/extensions/"+id, e, nil)
+	return c.UpdateExtensionWithContext(context.Background(), id, e)
+}
+
+// UpdateExtensionWithContext updates an extension by its ID.
+func (c *Client) UpdateExtensionWithContext(ctx context.Context, id string, e *Extension) (*Extension, error) {
+	resp, err := c.put(ctx, "/extensions/"+id, e, nil)
 	return getExtensionFromResponse(c, resp, err)
 }
 
@@ -75,7 +114,7 @@ func getExtensionFromResponse(c *Client, resp *http.Response, err error) (*Exten
 		return nil, fmt.Errorf("Could not decode JSON response: %v", dErr)
 	}
 
-	rootNode := "extension"
+	const rootNode = "extension"
 
 	t, nodeOK := target[rootNode]
 	if !nodeOK {