From 708c1758e4c021b530a331cc365f61b659248eb7 Mon Sep 17 00:00:00 2001
From: Letian Zhang <zlthooray@gmail.com>
Date: Fri, 4 Jan 2019 20:22:07 +0800
Subject: [PATCH] point to the same context after http.Request created

---
 context18_test.go | 23 +++++++++++++++++++++++
 request17.go      |  6 ++++++
 2 files changed, 29 insertions(+)

diff --git a/context18_test.go b/context18_test.go
index 93159b75..b984b709 100644
--- a/context18_test.go
+++ b/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 {
diff --git a/request17.go b/request17.go
index 3779673c..59845ba8 100644
--- a/request17.go
+++ b/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
 }