point to the same context after http.Request created

Author: clippit

context18_test.go

@@ -10,6 +10,7 @@ package resty
 
 import (
 	"context"
+	"net/http"
 	"net/url"
 	"testing"
 )
@@ -22,6 +23,28 @@ func TestRequestContext(t *testing.T) {
 	assertNotNil(t, r.Context())
 }
 
+func TestContextWithPreRequestHook(t *testing.T) {
+	ts := createGetServer(t)
+	defer ts.Close()
+
+	c := dc()
+	c.SetPreRequestHook(func(cl *Client, r *Request) error {
+		type ctxKey int
+		var key ctxKey
+		val := "test-value"
+		ctx := context.WithValue(r.Context(), key, val)
+		r.SetContext(ctx)
+		ctxValue := r.RawRequest.Context().Value(key).(string)
+		assertEqual(t, val, ctxValue)
+		return nil
+	})
+
+	resp, err := c.R().Get(ts.URL)
+
+	assertError(t, err)
+	assertEqual(t, http.StatusOK, resp.StatusCode())
+}
+
 func errIsContextCanceled(err error) bool {
 	ue, ok := err.(*url.Error)
 	if !ok {

request17.go

@@ -55,6 +55,9 @@ type Request struct {
 // Context method returns the Context if its already set in request
 // otherwise it creates new one using `context.Background()`.
 func (r *Request) Context() context.Context {
+	if r.RawRequest != nil {
+		return r.RawRequest.Context()
+	}
 	if r.ctx == nil {
 		return context.Background()
 	}
@@ -67,6 +70,9 @@ func (r *Request) Context() context.Context {
 // documentation.
 func (r *Request) SetContext(ctx context.Context) *Request {
 	r.ctx = ctx
+	if r.RawRequest != nil {
+		r.addContextIfAvailable()
+	}
 	return r
 }