Merge pull request #12 from wklken/ft_context

Ft context

gorequest.go

@@ -3,6 +3,7 @@ package gorequest
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
@@ -62,6 +63,7 @@ type SuperAgent struct {
 	Retryable            superAgentRetryable
 	DoNotClearSuperAgent bool
 	isClone              bool
+	ctx                  context.Context
 }
 
 var DisableTransportSwap = false
@@ -94,6 +96,7 @@ func New() *SuperAgent {
 		CurlCommand:       false,
 		logger:            log.New(os.Stderr, "[gorequest]", log.LstdFlags),
 		isClone:           false,
+		ctx:               nil,
 	}
 	// disable keep alives by default, see this issue https://github.com/parnurzeal/gorequest/issues/75
 	s.Transport.DisableKeepAlives = true
@@ -144,10 +147,16 @@ func (s *SuperAgent) Clone() *SuperAgent {
 		Retryable:            copyRetryable(s.Retryable),
 		DoNotClearSuperAgent: true,
 		isClone:              true,
+		ctx:                  s.ctx,
 	}
 	return clone
 }
 
+func (s *SuperAgent) Context(ctx context.Context) *SuperAgent {
+	s.ctx = ctx
+	return s
+}
+
 // SetDebug enable the debug mode which logs request/response detail.
 func (s *SuperAgent) SetDebug(enable bool) *SuperAgent {
 	s.Debug = enable
@@ -191,6 +200,7 @@ func (s *SuperAgent) ClearSuperAgent() {
 	s.TargetType = TypeJSON
 	s.Cookies = make([]*http.Cookie, 0)
 	s.Errors = nil
+	s.ctx = nil
 }
 
 // CustomMethod is just a wrapper to initialize SuperAgent instance by method string.
@@ -1291,6 +1301,11 @@ func (s *SuperAgent) MakeRequest() (*http.Request, error) {
 	if req, err = http.NewRequest(s.Method, s.Url, contentReader); err != nil {
 		return nil, err
 	}
+
+	if s.ctx != nil {
+		req = req.WithContext(s.ctx)
+	}
+
 	for k, vals := range s.Header {
 		for _, v := range vals {
 			req.Header.Add(k, v)

gorequest_test.go

@@ -2,6 +2,7 @@ package gorequest
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
@@ -2607,3 +2608,24 @@ func TestSetDebugByEnvironmentVar(t *testing.T) {
 		t.Fatalf("\nExpected gorequest not to log request and response object if GOREQUEST_DEBUG is not set.")
 	}
 }
+
+func TestContext(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// check method is GET before going to check other features
+		time.Sleep(1 * time.Second)
+
+		fmt.Fprintln(w, "Hello, client")
+	}))
+
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
+	defer cancel()
+
+	_, _, errs := New().Get(ts.URL).Context(ctx).EndBytes()
+	if len(errs) == 0 {
+		t.Fatalf("Expected error, got no error")
+	}
+
+	if !strings.HasSuffix(errs[0].Error(), "context deadline exceeded") {
+		t.Fatalf("Expected context deadline exceeded error, got %v", errs[0].Error())
+	}
+}