internal: Use provided context in subsequent request

Currently the HTTP request does not set the given context.
This change sets the context (if not nil) on the request.

Change-Id: I4bb21636d05050a68ba70ce92f9bf9ba608fbfad
Reviewed-on: https://go-review.googlesource.com/45370
Run-TryBot: Chris Broadfoot <cbro@golang.org>
Reviewed-by: Jaana Burcu Dogan <jbd@google.com>
Reviewed-by: Chris Broadfoot <cbro@golang.org>

2 files changed, 26 insertions, 1 deletions

diff --git a/internal/token.go b/internal/token.go
index efe7af5..766cfb6 100644
--- a/internal/token.go
+++ b/internal/token.go
@@ -18,6 +18,7 @@ import (
 	"time"
 
 	"golang.org/x/net/context"
+	"golang.org/x/net/context/ctxhttp"
 )
 
 // Token represents the crendentials used to authorize
@@ -189,7 +190,7 @@ func RetrieveToken(ctx context.Context, clientID, clientSecret, tokenURL string,
 	if !bustedAuth {
 		req.SetBasicAuth(clientID, clientSecret)
 	}
-	r, err := hc.Do(req)
+	r, err := ctxhttp.Do(ctx, hc, req)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/token_test.go b/internal/token_test.go
index 882de11..c19a462 100644
--- a/internal/token_test.go
+++ b/internal/token_test.go
@@ -79,3 +79,27 @@ func TestProviderAuthHeaderWorksDomain(t *testing.T) {
 		}
 	}
 }
+
+func TestRetrieveTokenWithContexts(t *testing.T) {
+	const clientID = "client-id"
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer ts.Close()
+
+	_, err := RetrieveToken(context.Background(), clientID, "", ts.URL, url.Values{})
+	if err != nil {
+		t.Errorf("RetrieveToken (with background context) = %v; want no error", err)
+	}
+
+	ctx, cancelfunc := context.WithCancel(context.Background())
+
+	cancellingts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		cancelfunc()
+	}))
+	defer cancellingts.Close()
+
+	_, err = RetrieveToken(ctx, clientID, "", cancellingts.URL, url.Values{})
+	if err == nil {
+		t.Errorf("RetrieveToken (with cancelled context) = nil; want error")
+	}
+}