Adding external cancelation context to the gentleman client

client.go

@@ -1,6 +1,7 @@
 package gentleman
 
 import (
+	gocontext "context"
 	"net/http"
 
 	"gopkg.in/h2non/gentleman.v2/context"
@@ -185,6 +186,17 @@ func (c *Client) CookieJar() *Client {
 	return c
 }
 
+// UseContext adds a cancelation context to the client to enable the use of early cancelation. This is useful for
+// server outgoing calls where we can attach the context from the incoming client. This will allow the downstream
+// calls to be canceled early on the case of a tcp close or http2 cancellation.
+// e.g.   client.Get().URL("someUrl").UseContext( incomingServerRequest.Context() ).Send()
+func (c *Client) UseContext(cancelContext gocontext.Context) *Client {
+	c.UseRequest(func(c *context.Context, handler context.Handler) {
+		handler.Next(c.SetCancelContext(cancelContext))
+	})
+	return c
+}
+
 // Use uses a new plugin to the middleware stack.
 func (c *Client) Use(p plugin.Plugin) *Client {
 	c.Middleware.Use(p)

client_test.go

@@ -1,13 +1,16 @@
 package gentleman
 
 import (
+	gocontext "context"
 	"errors"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/nbio/st"
+	"github.com/stretchr/testify/assert"
+
 	"gopkg.in/h2non/gentleman.v2/context"
 )
 
@@ -53,7 +56,7 @@ func TestClientRequestMiddleware(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Server", r.Header.Get("Client"))
 		w.Header().Set("Agent", r.Header.Get("Agent"))
-		fmt.Fprintln(w, "Hello, world")
+		_, _ = fmt.Fprintln(w, "Hello, world")
 	}))
 	defer ts.Close()
 
@@ -79,7 +82,7 @@ func TestClientRequestMiddleware(t *testing.T) {
 
 func TestClientRequestResponseMiddleware(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		fmt.Fprintln(w, "Hello, world")
+		_, _ = fmt.Fprintln(w, "Hello, world")
 	}))
 	defer ts.Close()
 
@@ -114,6 +117,8 @@ func TestClientErrorMiddleware(t *testing.T) {
 
 	req := client.Request()
 	res, err := req.Do()
+	assert.Error(t, err)
+	assert.NotNil(t, res)
 	st.Expect(t, err.Error(), "error: foo error")
 	st.Expect(t, res.Ok, false)
 	st.Expect(t, res.StatusCode, 0)
@@ -132,6 +137,9 @@ func TestClientCustomPhaseMiddleware(t *testing.T) {
 
 	req := client.Request()
 	res, err := req.Do()
+
+	assert.Error(t, err)
+	assert.NotNil(t, res)
 	st.Expect(t, err.Error(), "error: foo error")
 	st.Expect(t, res.Ok, false)
 	st.Expect(t, res.StatusCode, 0)
@@ -299,3 +307,17 @@ func TestClientVerbMethods(t *testing.T) {
 		t.Errorf("Invalid request method: %s", req.Context.Request.Method)
 	}
 }
+
+func TestClientWithCanceledContext(t *testing.T) {
+
+	ctx, cancel := gocontext.WithCancel(gocontext.Background())
+	cancel()
+	_, err := New().
+		URL("http://localhost:8999").
+		UseContext(ctx).
+		Post().
+		Path("/test").
+		Send()
+	assert.EqualError(t, err, "Post \"http://localhost:8999/test\": context canceled")
+
+}

context/context.go

@@ -304,6 +304,12 @@ func (c *Context) Value(key interface{}) interface{} {
 	return c.Request.Context().Value(key)
 }
 
+func (c *Context) SetCancelContext(ctx context.Context) *Context {
+	golRequestContext := context.WithValue(ctx, Key, c.Value(Key))
+	c.Request = c.Request.WithContext(golRequestContext)
+	return c
+}
+
 // emptyContext creates a new empty context.Context
 func emptyContext() context.Context {
 	return context.WithValue(context.Background(), Key, Store{})

plugins/proxy/proxy_test.go

@@ -2,6 +2,8 @@ package proxy
 
 import (
 	"github.com/nbio/st"
+	"github.com/stretchr/testify/assert"
+
 	"gopkg.in/h2non/gentleman.v2/context"
 	"net/http"
 	"testing"
@@ -37,8 +39,8 @@ func TestProxyParseError(t *testing.T) {
 
 	transport := ctx.Client.Transport.(*http.Transport)
 	_, err := transport.Proxy(ctx.Request)
-
-	st.Expect(t, err.Error(), "parse ://: missing protocol scheme")
+	assert.Error(t, err)
+	st.Expect(t, err.Error(), "parse \"://\": missing protocol scheme")
 }
 
 type handler struct {